From cd6d44bbaf8eb924fa347f8ef17107e2d4014dbb Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 8 Jan 2024 20:40:50 +0000 Subject: [PATCH 01/19] Add gitignores for metals (#2) --- .gitignore | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 600ab55..037ae54 100644 --- a/.gitignore +++ b/.gitignore @@ -2,8 +2,11 @@ .idea *.iml .vscode +.metals +.bloop # Scala related files +scala/project scala/target scala/project/target/ scala/project/project/target/ @@ -16,4 +19,4 @@ scala/src/test/scala/Playground.scala **/dist # Mac files -.DS_Store \ No newline at end of file +.DS_Store From 3f84a30c211cb0d06ea32ab63c48214e481917a3 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 8 Jan 2024 21:18:47 +0000 Subject: [PATCH 02/19] Update for spark 3.5.0 (#3) * Working with Spark 3.5.0 * Use SparkSession instead of deprecated SQLContext on the python side * Correction in init * Update tests for new spark session configuration --- python/ackuq/pit/context.py | 10 +++++----- python/tests/utils.py | 3 +-- scala/build.sbt | 4 ++-- scala/src/main/scala/EarlyStopSortMerge.scala | 2 +- scala/src/main/scala/execution/PITJoinExec.scala | 2 +- 5 files changed, 10 insertions(+), 11 deletions(-) diff --git a/python/ackuq/pit/context.py b/python/ackuq/pit/context.py index 000f3c4..862c916 100644 --- a/python/ackuq/pit/context.py +++ b/python/ackuq/pit/context.py @@ -25,7 +25,7 @@ from typing import Any, List, Optional, Tuple from py4j.java_gateway import JavaPackage -from pyspark.sql import Column, DataFrame, SQLContext +from pyspark.sql import Column, DataFrame, SparkSession from pyspark.sql.column import _to_java_column, _to_seq # type: ignore from pyspark.sql.functions import lit @@ -39,7 +39,7 @@ class PitContext(object): def _init_early_stop_sort_merge(self) -> None: # Call the init function in the Scala object, # will register the strategies and rules - self._essm.init(self._sql_context.sparkSession._jsparkSession) # type: ignore + self._essm.init(self._spark._jsparkSession) # type: ignore CLASSPATH_ERROR_MSG = ( "Java class {} could not be imported, check that it is included in the JVM" @@ -67,9 +67,9 @@ def _check_classpath(self): self.CLASSPATH_ERROR_MSG.format("io.github.ackuq.pit.Exploding") ) - def __init__(self, sql_context: SQLContext) -> None: - self._sql_context = sql_context - self._sc = self._sql_context._sc # type: ignore + def __init__(self, spark: SparkSession) -> None: + self._spark = spark + self._sc = self._spark.sparkContext self._jsc = self._sc._jsc self._jvm = self._sc._jvm diff --git a/python/tests/utils.py b/python/tests/utils.py index 0387822..5b6b3bb 100644 --- a/python/tests/utils.py +++ b/python/tests/utils.py @@ -44,8 +44,7 @@ def setUp(self) -> None: .config("spark.sql.shuffle.partitions", 1) .getOrCreate() ) - self.sql_context = SQLContext(self.spark.sparkContext) - self.pit_context = PitContext(self.sql_context) + self.pit_context = PitContext(self.spark) def tearDown(self) -> None: self.spark.stop() diff --git a/scala/build.sbt b/scala/build.sbt index 74a51ac..d034eec 100644 --- a/scala/build.sbt +++ b/scala/build.sbt @@ -12,8 +12,8 @@ lazy val root = (project in file(".")) ) libraryDependencies ++= Seq( - "org.apache.spark" %% "spark-core" % "3.2.1" % "provided", - "org.apache.spark" %% "spark-sql" % "3.2.1" % "provided", + "org.apache.spark" %% "spark-core" % "3.5.0" % "provided", + "org.apache.spark" %% "spark-sql" % "3.5.0" % "provided", "org.scalactic" %% "scalactic" % "3.2.12", "org.scalatest" %% "scalatest" % "3.2.12" % "test" ) diff --git a/scala/src/main/scala/EarlyStopSortMerge.scala b/scala/src/main/scala/EarlyStopSortMerge.scala index 5f439e5..6853bfd 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -48,7 +48,7 @@ object EarlyStopSortMerge { spark.experimental.extraStrategies = spark.experimental.extraStrategies :+ CustomStrategy } - if (!spark.experimental.extraStrategies.contains(PITRule)) { + if (!spark.experimental.extraOptimizations.contains(PITRule)) { spark.experimental.extraOptimizations = spark.experimental.extraOptimizations :+ PITRule } diff --git a/scala/src/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index c3f14cc..1546fa0 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -75,7 +75,7 @@ protected[pit] case class PITJoinExec( // TODO: This should be improved, but for now just keep everything in one partition AllTuples :: AllTuples :: Nil } else { - HashClusteredDistribution(leftEquiKeys) :: HashClusteredDistribution( + ClusteredDistribution(leftEquiKeys) :: ClusteredDistribution( rightEquiKeys ) :: Nil } From e7be94134f193e286cd1153c09bba7c444548f76 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Thu, 11 Jan 2024 18:32:07 +0000 Subject: [PATCH 03/19] Fix correctness bugs (#5) * Set spark log level to error * Run all tests with and without codegen * Re-name * Update tests to include product of all the things that were originally tested * Tests with no backward rows to match with * Add tests for empty dataframes * Add tests for preserving nulls in the input * Add tests for preserving nulls when tolerance is used * Fix codegen__left_join_some_rows_have_no_backwards_match * Fix interpreted__left_join_right_dataframe_is_empty * Fix interpretted__inner_join_searching_backward_for_matches_with_tolerance * Tidy * Remove debug `.show()`s --- .../main/scala/execution/PITJoinExec.scala | 22 +- .../src/test/scala/EarlyStopMergeTests.scala | 368 +++++++++++++++--- .../test/scala/SparkSessionTestWrapper.scala | 2 + scala/src/test/scala/data/SmallData.scala | 92 +++-- .../test/scala/data/SmallDataSortMerge.scala | 146 ++++++- 5 files changed, 512 insertions(+), 118 deletions(-) diff --git a/scala/src/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index 1546fa0..77bb788 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -202,14 +202,16 @@ protected[pit] case class PITJoinExec( new GenericInternalRow(right.output.length) override def advanceNext(): Boolean = { - while (smjScanner.findNextInnerJoinRows()) { + while (smjScanner.findNextLeftOuterJoinRows()) { currentRightMatch = smjScanner.getBufferedMatch currentLeftRow = smjScanner.getStreamedRow - if (returnNulls && currentRightMatch == null) { - joinRow(currentLeftRow, nullRightRow) - joinRow.withRight(nullRightRow) - numOutputRows += 1 - return true + if (currentRightMatch == null) { + if (returnNulls) { + joinRow(currentLeftRow, nullRightRow) + joinRow.withRight(nullRightRow) + numOutputRows += 1 + return true + } } else { joinRow(currentLeftRow, currentRightMatch) if (boundCondition(joinRow)) { @@ -493,6 +495,7 @@ protected[pit] case class PITJoinExec( | do { | if($rightRow == null) { | if(!rightIter.hasNext()) { + | $matched = null; | return false; | } | $rightRow = (InternalRow) rightIter.next(); @@ -745,7 +748,7 @@ protected[pit] class PITJoinScanner( * returns true, then [[getStreamedRow]] can be called to construct the * joinresults. */ - final def findNextInnerJoinRows(): Boolean = { + final def findNextLeftOuterJoinRows(): Boolean = { while ( advancedStreamed() && streamedRowEquiKey.anyNull && streamedRowPITKey.anyNull ) { @@ -777,7 +780,10 @@ protected[pit] class PITJoinScanner( matchJoinEquiKey = null matchJoinPITKey = null bufferedMatch = null - false + // If not returnNulls then we can exit early without searching through the remainder of the + // streamed row. If returnNulls then it doesn't matter that there are no more candidate rows + // to match with - we still need to return the streamed row so they can be matched with nulls. + returnNulls } else { // Advance both the streamed and buffered iterators to find the next pair of matching rows. var equiComp = diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index 29a0c05..b9271c1 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -24,51 +24,320 @@ package io.github.ackuq.pit -import EarlyStopSortMerge.pit -import data.SmallDataSortMerge - +import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.StructType import org.scalatest.flatspec.AnyFlatSpec +import org.apache.spark.sql.DataFrame +import EarlyStopSortMerge.pit +import data.SmallDataSortMerge class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { EarlyStopSortMerge.init(spark) val smallData = new SmallDataSortMerge(spark) - it should "Perform a PIT join with two dataframes, aligned timestamps" in { - val fg1 = smallData.fg1 - val fg2 = smallData.fg2 + def testBothCodegenAndInterpreted(name: String)(f: => Unit): Unit = { + it should (s"codegen__$name") in { + spark.conf.set("spark.sql.codegen.wholeStage", "true") + spark.conf.set("spark.sql.codegen.factoryMode", "CODEGEN_ONLY") + f + } + + it should (s"interpreted__$name") in { + spark.conf.set("spark.sql.codegen.wholeStage", "false") + spark.conf.set("spark.sql.codegen.factoryMode", "NO_CODEGEN") + f + } + } + + def testSearchingBackwardForMatches( + joinType: String, + leftDataFrame: DataFrame, + rightDataFrame: DataFrame, + expectedDataFrame: DataFrame, + expectedSchema: StructType, + tolerance: Int + ): Unit = { val pitJoin = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") + leftDataFrame.join( + rightDataFrame, + pit( + leftDataFrame("ts"), + rightDataFrame("ts"), + lit(tolerance) + ) && leftDataFrame("id") === rightDataFrame("id"), + joinType ) - assert(!pitJoin.isEmpty) - // Assert same schema - assert(pitJoin.schema.equals(smallData.PIT_1_2.schema)) - // Assert same elements - assert(pitJoin.collect().sameElements(smallData.PIT_1_2.collect())) + assert(pitJoin.schema.equals(expectedSchema)) + assert(pitJoin.collect().sameElements(expectedDataFrame.collect())) } - it should "Perform a PIT join with two dataframes, misaligned timestamps" in { - val fg1 = smallData.fg1 - val fg2 = smallData.fg3 + testBothCodegenAndInterpreted( + "inner_join_with_aligned_timestamps_no_tolerance" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.fg2, + smallData.PIT_1_2, + smallData.PIT_2_schema, + 0 + ) + } - val pitJoin = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") - ) + testBothCodegenAndInterpreted( + "left_join_with_aligned_timestamps_no_tolerance" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1, + smallData.fg2, + smallData.PIT_1_2, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_with_aligned_timestamps_with_tolerance" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.fg2, + smallData.PIT_1_2, + smallData.PIT_2_schema, + 1 + ) + } + + testBothCodegenAndInterpreted( + "left_join_with_aligned_timestamps_with_tolerance" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1, + smallData.fg2, + smallData.PIT_1_2, + smallData.PIT_2_OUTER_schema, + 1 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_searching_backward_for_matches_no_tolerance" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.fg3, + smallData.PIT_1_3, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_searching_backward_for_matches_no_tolerance" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1, + smallData.fg3, + smallData.PIT_1_3, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_searching_backward_for_matches_with_tolerance" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.fg3, + smallData.PIT_1_3_T1, + smallData.PIT_2_schema, + 1 + ) + } + + testBothCodegenAndInterpreted( + "left_join_searching_backward_for_matches_with_tolerance" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1, + smallData.fg3, + smallData.PIT_1_3_T1_OUTER, + smallData.PIT_2_OUTER_schema, + 1 + ) + } - assert(!pitJoin.isEmpty) - // Assert same schema - assert(pitJoin.schema.equals(smallData.PIT_1_3.schema)) - // Assert same elements - assert(pitJoin.collect().sameElements(smallData.PIT_1_3.collect())) + testBothCodegenAndInterpreted( + "inner_join_some_rows_have_no_backwards_match" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg3, + smallData.fg1, + smallData.PIT_3_1, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_some_rows_have_no_backwards_match" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg3, + smallData.fg1, + smallData.PIT_3_1_OUTER, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_right_dataframe_is_empty" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.empty, + smallData.PIT_EMPTY, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_right_dataframe_is_empty" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1, + smallData.empty, + smallData.PIT_1_EMPTY_OUTER, + smallData.PIT_2_OUTER_schema, + 0 + ) } - it should "Perform a PIT join with three dataframes, misaligned timestamps" in { + testBothCodegenAndInterpreted( + "inner_join_left_dataframe_is_empty" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1, + smallData.empty, + smallData.PIT_EMPTY, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_left_dataframe_is_empty" + ) { + testSearchingBackwardForMatches( + "left", + smallData.empty, + smallData.fg1, + smallData.PIT_EMPTY, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_both_dataframes_are_empty" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.empty, + smallData.empty2, // Don't want to test self join + smallData.PIT_EMPTY, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_both_dataframes_are_empty" + ) { + testSearchingBackwardForMatches( + "left", + smallData.empty, + smallData.empty2, // Don't want to test self join + smallData.PIT_EMPTY, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_preserves_input_nulls" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg3_with_nulls, + smallData.fg1_with_nulls, + smallData.PIT_3_1_WITH_NULLS, + smallData.PIT_2_NULLABLE_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_preserves_input_nulls" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg3_with_nulls, + smallData.fg1_with_nulls, + smallData.PIT_3_1_WITH_NULLS_OUTER, + smallData.PIT_2_NULLABLE_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_preserves_input_nulls_with_tolerance" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg3_with_nulls, + smallData.fg1_with_nulls, + smallData.PIT_3_1_T1_WITH_NULLS, + smallData.PIT_2_NULLABLE_schema, + 1 + ) + } + + testBothCodegenAndInterpreted( + "left_join_preserves_input_nulls_with_tolerance" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg3_with_nulls, + smallData.fg1_with_nulls, + smallData.PIT_3_1_T1_WITH_NULLS_OUTER, + smallData.PIT_2_NULLABLE_OUTER_schema, + 1 + ) + } + + def testJoiningThreeDataframes( + joinType: String, + expectedSchema: StructType + ): Unit = { val fg1 = smallData.fg1 val fg2 = smallData.fg2 val fg3 = smallData.fg3 @@ -76,50 +345,25 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { val left = fg1.join( fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") + pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id"), + joinType ) val pitJoin = left.join( fg3, - pit(fg1("ts"), fg3("ts"), lit(0)) && fg1("id") === fg3("id") + pit(fg1("ts"), fg3("ts"), lit(0)) && fg1("id") === fg3("id"), + joinType ) - assert(!pitJoin.isEmpty) - // Assert same schema - assert(pitJoin.schema.equals(smallData.PIT_1_2_3.schema)) - // Assert same elements + assert(pitJoin.schema.equals(expectedSchema)) assert(pitJoin.collect().sameElements(smallData.PIT_1_2_3.collect())) } - it should "Be able to perform a PIT join with tolerance, misaligned timestamps" in { - val fg1 = smallData.fg1 - val fg2 = smallData.fg3 - - val pitJoin = fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(1)) && fg1("id") === fg2("id") - ) - assert(!pitJoin.isEmpty) - // Assert same schema - assert(pitJoin.schema.equals(smallData.PIT_1_3_T1.schema)) - // Assert same elements - assert(pitJoin.collect().sameElements(smallData.PIT_1_3_T1.collect())) + testBothCodegenAndInterpreted("inner_join_three_dataframes") { + testJoiningThreeDataframes("inner", smallData.PIT_3_schema) } - - it should "Be able to perform a left outer PIT join with tolerance, misaligned timestamps" in { - val fg1 = smallData.fg1 - val fg2 = smallData.fg3 - - val pitJoin = fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(1)) && fg1("id") === fg2("id"), - "left" - ) - assert(!pitJoin.isEmpty) - // Assert same schema - assert(pitJoin.schema.equals(smallData.PIT_1_3_T1_OUTER.schema)) - // Assert same elements - assert(pitJoin.collect().sameElements(smallData.PIT_1_3_T1_OUTER.collect())) + testBothCodegenAndInterpreted("left_join_three_dataframes") { + testJoiningThreeDataframes("left", smallData.PIT_3_OUTER_schema) } } diff --git a/scala/src/test/scala/SparkSessionTestWrapper.scala b/scala/src/test/scala/SparkSessionTestWrapper.scala index 9bf7fe0..51bbcfd 100644 --- a/scala/src/test/scala/SparkSessionTestWrapper.scala +++ b/scala/src/test/scala/SparkSessionTestWrapper.scala @@ -34,4 +34,6 @@ trait SparkSessionTestWrapper { .config("spark.ui.showConsoleProgress", value = false) .config("spark.sql.shuffle.partitions", 1) .getOrCreate() + + spark.sparkContext.setLogLevel("ERROR") } diff --git a/scala/src/test/scala/data/SmallData.scala b/scala/src/test/scala/data/SmallData.scala index 8dae94d..5905c3b 100644 --- a/scala/src/test/scala/data/SmallData.scala +++ b/scala/src/test/scala/data/SmallData.scala @@ -25,50 +25,78 @@ package io.github.ackuq.pit package data -import org.apache.spark.sql.types.{ - IntegerType, - StringType, - StructField, - StructType -} -import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType class SmallData(spark: SparkSession) { - private val DATA_RAW = Seq( - Seq( - Row(1, 4, "1z"), - Row(1, 5, "1x"), - Row(2, 6, "2x"), - Row(1, 7, "1y"), - Row(2, 8, "2y") - ), - Seq( - Row(1, 4, "1z"), - Row(1, 5, "1x"), - Row(2, 6, "2x"), - Row(1, 7, "1y"), - Row(2, 8, "2y") - ), + val RAW_1 = Seq( + Row(1, 4, "1z"), + Row(1, 5, "1x"), + Row(2, 6, "2x"), + Row(1, 7, "1y"), + Row(2, 8, "2y") + ) + val RAW_1_WITH_NULLS = Seq( + Row(1, 4, "1z"), + Row(1, 5, null), + Row(2, 6, "2x"), + Row(1, 7, "1y"), + Row(2, 8, "2y") + ) + val RAW_3 = Seq( + Row(1, 1, "f3-1-1"), + Row(2, 2, "f3-2-2"), + Row(1, 6, "f3-1-6"), + Row(2, 8, "f3-2-8"), + Row(1, 10, "f3-1-10") + ) + val RAW_3_WITH_NULLS = Seq( + Row(1, 1, "f3-1-1"), + Row(2, 2, "f3-2-2"), + Row(1, 6, "f3-1-6"), + Row(2, 8, null), + Row(1, 10, "f3-1-10") + ) + val schema: StructType = StructType( Seq( - Row(1, 1, "f3-1-1"), - Row(2, 2, "f3-2-2"), - Row(1, 6, "f3-1-6"), - Row(2, 8, "f3-2-8"), - Row(1, 10, "f3-1-10") + StructField("id", IntegerType, nullable = false), + StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = false) ) ) - private val schema: StructType = StructType( + val schema_nullable: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), - StructField("value", StringType, nullable = false) + StructField("value", StringType, nullable = true) ) ) val fg1: DataFrame = - spark.createDataFrame(spark.sparkContext.parallelize(DATA_RAW.head), schema) + spark.createDataFrame(spark.sparkContext.parallelize(RAW_1), schema) + val fg1_with_nulls: DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(RAW_1_WITH_NULLS), + schema_nullable + ) val fg2: DataFrame = - spark.createDataFrame(spark.sparkContext.parallelize(DATA_RAW(1)), schema) + spark.createDataFrame(spark.sparkContext.parallelize(RAW_1), schema) val fg3: DataFrame = - spark.createDataFrame(spark.sparkContext.parallelize(DATA_RAW(2)), schema) + spark.createDataFrame(spark.sparkContext.parallelize(RAW_3), schema) + val fg3_with_nulls: DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(RAW_3_WITH_NULLS), + schema_nullable + ) + + val empty: DataFrame = + spark.createDataFrame(spark.sparkContext.emptyRDD[Row], schema) + + val empty2: DataFrame = + spark.createDataFrame(spark.sparkContext.emptyRDD[Row], schema) } diff --git a/scala/src/test/scala/data/SmallDataSortMerge.scala b/scala/src/test/scala/data/SmallDataSortMerge.scala index 07d80d2..aa2deb5 100644 --- a/scala/src/test/scala/data/SmallDataSortMerge.scala +++ b/scala/src/test/scala/data/SmallDataSortMerge.scala @@ -25,49 +25,90 @@ package io.github.ackuq.pit package data -import org.apache.spark.sql.types.{ - IntegerType, - StringType, - StructField, - StructType -} -import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { - - private val PIT_1_2_RAW = Seq( + val PIT_1_2_RAW = Seq( Row(2, 8, "2y", 2, 8, "2y"), Row(2, 6, "2x", 2, 6, "2x"), Row(1, 7, "1y", 1, 7, "1y"), Row(1, 5, "1x", 1, 5, "1x"), Row(1, 4, "1z", 1, 4, "1z") ) - private val PIT_1_3_RAW = Seq( + val PIT_1_EMPTY_RAW = Seq( + Row(2, 8, "2y", null, null, null), + Row(2, 6, "2x", null, null, null), + Row(1, 7, "1y", null, null, null), + Row(1, 5, "1x", null, null, null), + Row(1, 4, "1z", null, null, null) + ) + val PIT_1_3_RAW = Seq( Row(2, 8, "2y", 2, 8, "f3-2-8"), Row(2, 6, "2x", 2, 2, "f3-2-2"), Row(1, 7, "1y", 1, 6, "f3-1-6"), Row(1, 5, "1x", 1, 1, "f3-1-1"), Row(1, 4, "1z", 1, 1, "f3-1-1") ) - private val PIT_1_3_T1_RAW = Seq( + val PIT_3_1_RAW = Seq( + Row(2, 8, "f3-2-8", 2, 8, "2y"), + Row(1, 10, "f3-1-10", 1, 7, "1y"), + Row(1, 6, "f3-1-6", 1, 5, "1x") + ) + val PIT_3_1_WITH_NULLS_RAW = Seq( + Row(2, 8, null, 2, 8, "2y"), + Row(1, 10, "f3-1-10", 1, 7, "1y"), + Row(1, 6, "f3-1-6", 1, 5, null) + ) + val PIT_3_1_T1_WITH_NULLS_RAW = Seq( + Row(2, 8, null, 2, 8, "2y"), + Row(1, 6, "f3-1-6", 1, 5, null) + ) + val PIT_3_1_OUTER_RAW = Seq( + Row(2, 8, "f3-2-8", 2, 8, "2y"), + Row(2, 2, "f3-2-2", null, null, null), + Row(1, 10, "f3-1-10", 1, 7, "1y"), + Row(1, 6, "f3-1-6", 1, 5, "1x"), + Row(1, 1, "f3-1-1", null, null, null) + ) + val PIT_3_1_T1_WITH_NULLS_OUTER_RAW = Seq( + Row(2, 8, null, 2, 8, "2y"), + Row(2, 2, "f3-2-2", null, null, null), + Row(1, 10, "f3-1-10", null, null, null), + Row(1, 6, "f3-1-6", 1, 5, null), + Row(1, 1, "f3-1-1", null, null, null) + ) + val PIT_3_1_WITH_NULLS_OUTER_RAW = Seq( + Row(2, 8, null, 2, 8, "2y"), + Row(2, 2, "f3-2-2", null, null, null), + Row(1, 10, "f3-1-10", 1, 7, "1y"), + Row(1, 6, "f3-1-6", 1, 5, null), + Row(1, 1, "f3-1-1", null, null, null) + ) + val PIT_1_3_T1_RAW = Seq( Row(2, 8, "2y", 2, 8, "f3-2-8"), Row(1, 7, "1y", 1, 6, "f3-1-6") ) - private val PIT_1_3_T1_OUTER_RAW = Seq( + val PIT_1_3_T1_OUTER_RAW = Seq( Row(2, 8, "2y", 2, 8, "f3-2-8"), Row(2, 6, "2x", null, null, null), Row(1, 7, "1y", 1, 6, "f3-1-6"), Row(1, 5, "1x", null, null, null), Row(1, 4, "1z", null, null, null) ) - private val PIT_1_2_3_RAW = Seq( + val PIT_1_2_3_RAW = Seq( Row(2, 8, "2y", 2, 8, "2y", 2, 8, "f3-2-8"), Row(2, 6, "2x", 2, 6, "2x", 2, 2, "f3-2-2"), Row(1, 7, "1y", 1, 7, "1y", 1, 6, "f3-1-6"), Row(1, 5, "1x", 1, 5, "1x", 1, 1, "f3-1-1"), Row(1, 4, "1z", 1, 4, "1z", 1, 1, "f3-1-1") ) - private val PIT_2_schema: StructType = StructType( + val PIT_2_schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), @@ -77,7 +118,17 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { StructField("value", StringType, nullable = false) ) ) - private val PIT_2_OUTER_schema: StructType = StructType( + val PIT_2_NULLABLE_schema: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = false), + StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = true), + StructField("id", IntegerType, nullable = false), + StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = true) + ) + ) + val PIT_2_OUTER_schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), @@ -87,8 +138,18 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { StructField("value", StringType, nullable = true) ) ) + val PIT_2_NULLABLE_OUTER_schema: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = false), + StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = true), + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = true) + ) + ) - private val PIT_3_schema: StructType = StructType( + val PIT_3_schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), @@ -101,6 +162,19 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { StructField("value", StringType, nullable = false) ) ) + val PIT_3_OUTER_schema: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = false), + StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = false), + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = true), + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = true) + ) + ) val PIT_1_2: DataFrame = spark.createDataFrame( spark.sparkContext.parallelize(PIT_1_2_RAW), PIT_2_schema @@ -120,8 +194,48 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { PIT_2_OUTER_schema ) + val PIT_3_1_OUTER: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_OUTER_RAW), + PIT_2_OUTER_schema + ) + + val PIT_3_1: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_RAW), + PIT_2_OUTER_schema + ) + + val PIT_3_1_WITH_NULLS: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_WITH_NULLS_RAW), + PIT_2_NULLABLE_schema + ) + + val PIT_3_1_T1_WITH_NULLS: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_T1_WITH_NULLS_RAW), + PIT_2_NULLABLE_schema + ) + + val PIT_3_1_WITH_NULLS_OUTER: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_WITH_NULLS_OUTER_RAW), + PIT_2_NULLABLE_OUTER_schema + ) + + val PIT_3_1_T1_WITH_NULLS_OUTER: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_T1_WITH_NULLS_OUTER_RAW), + PIT_2_NULLABLE_OUTER_schema + ) + val PIT_1_2_3: DataFrame = spark.createDataFrame( spark.sparkContext.parallelize(PIT_1_2_3_RAW), PIT_3_schema ) + + val PIT_EMPTY: DataFrame = spark.createDataFrame( + spark.sparkContext.emptyRDD[Row], + PIT_2_schema + ) + + val PIT_1_EMPTY_OUTER: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_1_EMPTY_RAW), + PIT_2_OUTER_schema + ) } From e5fcee706517d064c8a9065a4eac20e51cd6f688 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 15 Jan 2024 13:48:05 +0000 Subject: [PATCH 04/19] Refactor interpreted code path + fix nulls in join keys (#6) * Start tests * Valid tests * Fix codegen version * Tests with duplicate rows * Delete unnecessary code. It was a legacy of the normal join's multi matching behaviour * Working inner joins with more code removed and tolerance applied inside main search loop * Working inner join * Tidy * Very minor clean ups * Remove some unnecessary branching * Fix schema assertion for left_join_duplicate_join_keys * Remove unneeded bound condition * Fix schema assertions for nulls in join keys * Update comments and rename variables * Auto-format * Remove slightly misleading comment * Remove completed TODO comment * More comment adjustments --- .../main/scala/execution/PITJoinExec.scala | 397 +++++++++--------- .../src/test/scala/EarlyStopMergeTests.scala | 83 +++- scala/src/test/scala/data/SmallData.scala | 78 +++- .../test/scala/data/SmallDataSortMerge.scala | 76 +++- 4 files changed, 412 insertions(+), 222 deletions(-) diff --git a/scala/src/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index 77bb788..fc750f3 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -20,19 +20,21 @@ package io.github.ackuq.pit package execution -import logical.{CustomJoinType, PITJoinType} - import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.JoinType import org.apache.spark.sql.catalyst.plans.physical._ -import org.apache.spark.sql.catalyst.plans.{Inner, JoinType} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.joins.ShuffledJoin -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.metric.SQLMetrics + +import logical.{CustomJoinType, PITJoinType} /** Performs a PIT join of two child relations. */ @@ -145,19 +147,7 @@ protected[pit] case class PITJoinExec( val numOutputRows = longMetric("numOutputRows") left.execute().zipPartitions(right.execute()) { (leftIter, rightIter) => - val boundCondition: InternalRow => Boolean = { - condition - .map { cond => - Predicate.create(cond, left.output ++ right.output).eval _ - - } - .getOrElse { (r: InternalRow) => - true - } - } - // An ordering that can be used to compare keys from both sides. - // Ordering of the equi-keys val equiOrder: Seq[SortOrder] = leftEquiKeys.map(_.dataType).zipWithIndex.map { case (dt, index) => @@ -202,27 +192,20 @@ protected[pit] case class PITJoinExec( new GenericInternalRow(right.output.length) override def advanceNext(): Boolean = { - while (smjScanner.findNextLeftOuterJoinRows()) { - currentRightMatch = smjScanner.getBufferedMatch - currentLeftRow = smjScanner.getStreamedRow + if (smjScanner.findNextJoinRow()) { + currentRightMatch = smjScanner.getRightMatch + currentLeftRow = smjScanner.getLeftRow if (currentRightMatch == null) { - if (returnNulls) { - joinRow(currentLeftRow, nullRightRow) - joinRow.withRight(nullRightRow) - numOutputRows += 1 - return true - } + joinRow(currentLeftRow, nullRightRow) + joinRow.withRight(nullRightRow) + numOutputRows += 1 + return true } else { joinRow(currentLeftRow, currentRightMatch) - if (boundCondition(joinRow)) { - numOutputRows += 1 - return true - } + numOutputRows += 1 + return true } } - - currentRightMatch = null - currentLeftRow = null false } @@ -438,10 +421,6 @@ protected[pit] case class PITJoinExec( } } - /** Generate a function to scan both left and right to find a match, returns - * the term for matched one row from left side and buffered rows from right - * side. - */ private def genScanner(ctx: CodegenContext) = { // Create class member for next row from both sides. // Inline mutable state since not many join operations in a task @@ -473,10 +452,32 @@ protected[pit] case class PITJoinExec( val matchedPITKeyVars = copyKeys(ctx, leftPITKeyVars, leftPitKeys) val matchedEquiKeyVars = copyKeys(ctx, leftEquiKeyVars, leftEquiKeys) + // Generate a function to scan both left and right sides to find a match. + // Return whether a match is found. + // + // `leftIter`: the iterator for left side. + // `rightIter`: the iterator for right side. + // `leftRow`: the current row from left side. + // When `leftIter` is empty, `leftRow` is null. + // `rightRow`: the current match candidate row from right side. + // `match`: the row from right side which matches the with `leftRow`. + // If there is no match with `leftRow`, `matches` is empty. + // + // The function has the following step: + // - Step 1: Find the next `leftRow` with non-null join keys. + // For `leftRow` with null join keys: + // 1. Inner join: skip this `leftRow` the row, and try the next one. + // 2. Left Outer join: keep the row and return false with null `match`. + // + // - Step 2: Find the first `match` from right side matching the join keys + // with `leftRow`. Return true after finding a match. + // For `leftRow` without `match`: + // 1. Inner join: skip this `leftRow` the row, and try the next one. + // 2. Left Outer join: keep the row and return false with null `match`. ctx.addNewFunction( - "findNextInnerJoinRows", + "findNextJoinRows", s""" - |private boolean findNextInnerJoinRows( + |private boolean findNextJoinRows( |scala.collection.Iterator leftIter, |scala.collection.Iterator rightIter |){ @@ -489,6 +490,10 @@ protected[pit] case class PITJoinExec( | ${leftPITKeyVars.map(_.code).mkString("\n")} | ${leftEquiKeyVars.map(_.code).mkString("\n")} | if($leftAnyNull) { + | if($returnNulls) { + | $matched = null; + | return false; + | } | $leftRow = null; | continue; | } @@ -638,7 +643,7 @@ protected[pit] case class PITJoinExec( if (returnNulls) { s""" |while($leftInput.hasNext()) { - | findNextInnerJoinRows($leftInput, $rightInput); + | findNextJoinRows($leftInput, $rightInput); | ${leftVarDecl.mkString("\n")} | ${beforeLoop.trim} | InternalRow $rightRow = (InternalRow) $matched; @@ -650,7 +655,7 @@ protected[pit] case class PITJoinExec( |""".stripMargin } else { s""" - |while (findNextInnerJoinRows($leftInput, $rightInput)) { + |while (findNextJoinRows($leftInput, $rightInput)) { | ${leftVarDecl.mkString("\n")} | ${beforeLoop.trim} | InternalRow $rightRow = (InternalRow) $matched; @@ -667,217 +672,217 @@ protected[pit] case class PITJoinExec( /** Helper class that is used to implement [[PITJoinExec]]. * - * To perform an inner (outer) join, users of this class call - * [[findNextInnerJoinRows()]] which returns `true` if a result has been - * produced and `false` otherwise. If a result has been produced, then the - * caller may call [[getStreamedRow]] to return the matching row from the - * streamed input For efficiency, both of these methods return mutable objects - * which are re-used across calls to the `findNext*JoinRows()` methods. - * - * @param streamedPITKeyGenerator - * a projection that produces PIT join keys from the streamed input. - * @param bufferedPITKeyGenerator - * a projection that produces PIT join keys from the buffered input. + * To perform an inner or outer join, users of this class call + * [[findNextJoinRow()]], which returns `true` if a result has been produced + * and `false` otherwise. If a result has been produced, then the caller may + * call [[getLeftRow]] and [[getRightMatch]] to get left and right sides + * respectively, of the joined row. Both of these methods return mutable + * objects which are re-used across calls to the `findNextJoinRow()`. + * @param leftPITKeyGenerator + * a projection that produces PIT join keys from the left input. + * @param rightPITKeyGenerator + * a projection that produces PIT join keys from the right input. * @param pitKeyOrdering * an ordering which can be used to compare PIT join keys. - * @param streamedEquiKeyGenerator - * a projection that produces join keys from the streamed input. - * @param bufferedEquiKeyGenerator - * a projection that produces join keys from the buffered input. + * @param leftEquiKeyGenerator + * a projection that produces join keys from the left input. + * @param rightEquiKeyGenerator + * a projection that produces join keys from the right input. * @param equiKeyOrdering * an ordering which can be used to compare equi join keys. - * @param streamedIter - * an input whose rows will be streamed. - * @param bufferedIter - * an input whose rows will be buffered to construct sequences of rows that - * have the same join key. + * @param leftIter + * an input whose rows will be left. + * @param rightIter + * an input whose rows will be right. * @param eagerCleanupResources * the eager cleanup function to be invoked when no join row found * @param returnNulls - * event if no PIT match found, return the left row with right columns filled - * with null + * if no PIT match found, return the left row with right columns filled with + * null (left outer join) * @param tolerance * tolerance for how long we want to look back, if value is 0, do no use * tolerance */ protected[pit] class PITJoinScanner( - streamedPITKeyGenerator: Projection, - bufferedPITKeyGenerator: Projection, + leftPITKeyGenerator: Projection, + rightPITKeyGenerator: Projection, pitKeyOrdering: Ordering[InternalRow], - streamedEquiKeyGenerator: Projection, - bufferedEquiKeyGenerator: Projection, + leftEquiKeyGenerator: Projection, + rightEquiKeyGenerator: Projection, equiKeyOrdering: Ordering[InternalRow], - streamedIter: RowIterator, - bufferedIter: RowIterator, + leftIter: RowIterator, + rightIter: RowIterator, eagerCleanupResources: () => Unit, returnNulls: Boolean = false, tolerance: Long = 0 ) { - private[this] var streamedRow: InternalRow = _ - private[this] var streamedRowEquiKey: InternalRow = _ - private[this] var streamedRowPITKey: InternalRow = _ - private[this] var bufferedRow: InternalRow = _ + private[this] var leftRow: InternalRow = _ + private[this] var leftRowEquiKey: InternalRow = _ + private[this] var leftRowPITKey: InternalRow = _ + private[this] var rightRow: InternalRow = _ // Note: this is guaranteed to never have any null columns: - private[this] var bufferedRowEquiKey: InternalRow = _ - private[this] var bufferedRowPITKey: InternalRow = _ - - /** The join key for the rows buffered in `bufferedMatches`, or null if - * `bufferedMatches` is empty - */ - private[this] var matchJoinEquiKey: InternalRow = _ - private[this] var matchJoinPITKey: InternalRow = _ + private[this] var rightRowEquiKey: InternalRow = _ + private[this] var rightRowPITKey: InternalRow = _ /** Contains the current match */ - private[this] var bufferedMatch: UnsafeRow = _ + private[this] var rightMatch: UnsafeRow = _ - // Initialization (note: do _not_ want to advance streamed here). - advancedBuffered() + // Initialization (note: do _not_ want to advance left here). + advancedRightToRowWithNullFreeJoinKeys() // --- Public methods --------------------------------------------------------------------------- - def getStreamedRow: InternalRow = streamedRow + def getLeftRow: InternalRow = leftRow - def getBufferedMatch: UnsafeRow = bufferedMatch + def getRightMatch: UnsafeRow = rightMatch - /** Advances both input iterators, stopping when we have found rows with - * matching join keys. If no join rows found, try to do the eager resources - * cleanup. - * + /** Advances left iterator by one then starts searching for a match. If + * returnNulls it will always return a result for every left row so it will + * only advance the right iterator when searching. If not returnNulls it will + * may advance both iterators when searching for a match. When there are no + * more potential matches it does eager cleanup. * @return * true if matching rows have been found and false otherwise. If this - * returns true, then [[getStreamedRow]] can be called to construct the - * joinresults. + * returns true, then [[getLeftRow]] and [[getRightMatch]] can be called to + * construct the join result. */ - final def findNextLeftOuterJoinRows(): Boolean = { - while ( - advancedStreamed() && streamedRowEquiKey.anyNull && streamedRowPITKey.anyNull - ) { - // Advance the streamed side of the join until we find the next row whose join key contains - // no nulls or we hit the end of the streamed iterator. - } - val found = if (streamedRow == null) { - // We have consumed the entire streamed iterator, so there can be no more matches. - matchJoinEquiKey = null - matchJoinPITKey = null - bufferedMatch = null - false - } else if ( - matchJoinEquiKey != null && equiKeyOrdering.compare( - streamedRowEquiKey, - matchJoinEquiKey - ) == 0 - && matchJoinPITKey != null && pitKeyOrdering.compare( - streamedRowPITKey, - matchJoinPITKey - // Streamed row key must be equal or greater than match - ) <= 0 - ) { - // The new streamed row has the same join key as the previous row, so return the same matches. - true - } else if (bufferedRow == null) { - // The streamed row's join key does not match the current batch of buffered rows and there are - // no more rows to read from the buffered iterator, so there can be no more matches. - matchJoinEquiKey = null - matchJoinPITKey = null - bufferedMatch = null - // If not returnNulls then we can exit early without searching through the remainder of the - // streamed row. If returnNulls then it doesn't matter that there are no more candidate rows - // to match with - we still need to return the streamed row so they can be matched with nulls. - returnNulls - } else { - // Advance both the streamed and buffered iterators to find the next pair of matching rows. - var equiComp = - equiKeyOrdering.compare(streamedRowEquiKey, bufferedRowEquiKey) - var pitComp = pitKeyOrdering.compare(streamedRowPITKey, bufferedRowPITKey) - do { - if (streamedRowPITKey.anyNull || streamedRowEquiKey.anyNull) { - advancedStreamed() + final def findNextJoinRow(): Boolean = { + // Advance the `leftRow` at the start of every call to avoid returning the same rows repeatedly. + val leftIteratorNotEmpty = advancedLeft() + val rightIteratorNotEmpty = rightRow != null + var keepSearchingPotentiallyMoreJoinRows: (Boolean, Boolean) = + ( + // If both iterators are non-empty then searching might find a match - keep searching. + leftIteratorNotEmpty && rightIteratorNotEmpty, + // For more join rows the left iterator must be non-empty and either returnNulls is allowed or + // the right iterator is non-empty. + leftIteratorNotEmpty && (returnNulls || rightIteratorNotEmpty) + ) + // Run the search. If not returnNulls this might require advancing both iterators. + while (keepSearchingPotentiallyMoreJoinRows._1) { + keepSearchingPotentiallyMoreJoinRows = { + if (leftRowPITKey.anyNull || leftRowEquiKey.anyNull) { + handleNoMatchForLeftRow() } else { - assert(!bufferedRowPITKey.anyNull && !bufferedRowEquiKey.anyNull) + assert(!rightRowPITKey.anyNull && !rightRowEquiKey.anyNull) - equiComp = - equiKeyOrdering.compare(streamedRowEquiKey, bufferedRowEquiKey) - pitComp = pitKeyOrdering.compare(streamedRowPITKey, bufferedRowPITKey) + val equiComp = + equiKeyOrdering.compare(leftRowEquiKey, rightRowEquiKey) + val pitComp = + pitKeyOrdering.compare(leftRowPITKey, rightRowPITKey) if (equiComp < 0) { - if (returnNulls) { - bufferedMatch = null - return true - } - advancedStreamed() - } else if (equiComp > 0) advancedBuffered() - else if (pitComp > 0) advancedBuffered() - } - } while (streamedRow != null && bufferedRow != null && (equiComp != 0 || pitComp > 0)) - if (returnNulls && streamedRow != null && bufferedRow == null) { - // Not a match found for the streamed row - bufferedMatch = null - true - } else if (streamedRow == null || bufferedRow == null) { - // We have either hit the end of one of the iterators, so there can be no more matches. - matchJoinEquiKey = null - matchJoinPITKey = null - bufferedMatch = null - false - } else { - // The streamed row's join key matches the current buffered row's join, only take this row - assert(equiComp == 0 && pitComp <= 0) - - if (tolerance > 0) { - val streamedTS = streamedRowPITKey.getLong(0) - val bufferedTS = bufferedRowPITKey.getLong(0) - if (streamedTS - bufferedTS > tolerance) { - // The one closest to the streamed row exceeds the tolerance, continue... - bufferedMatch = null - return true + // leftRowEquiKey > rightRowEquiKey + // Advance (decrement) `leftRow` to find next potential matches. + handleNoMatchForLeftRow() + } else if (equiComp > 0 || pitComp > 0) { + // leftRowEquiKey < rightRowEquiKey || leftRowPITKey < rightRowPITKey + // Advance (decrement) `rightRow` to find next potential matches. + handleRightRowDoesNotMatch() + } else if ( + tolerance > 0 && ( + leftRowPITKey + .getLong(0) - rightRowPITKey.getLong(0) > tolerance + ) + ) { + // Tolerance is enabled and `rightRow` is outside the tolerance. No other + // `rightRow` will be any closer to this `leftRow`, so advance to the + // next `leftRow`. + handleNoMatchForLeftRow() + } else { + // The left row's join key matches the current right row's join, only take this row + assert(equiComp == 0 && pitComp <= 0) + rightMatch = rightRow.asInstanceOf[UnsafeRow] + ( + false, // Stop searching after finding a valid match. + true // Potentially there will be another left row, therefore more join rows. + ) } } - - bufferedMatch = bufferedRow.asInstanceOf[UnsafeRow] - true } } - if (!found) eagerCleanupResources() - found + if (!keepSearchingPotentiallyMoreJoinRows._2) eagerCleanupResources() + keepSearchingPotentiallyMoreJoinRows._2 } // --- Private methods -------------------------------------------------------------------------- - /** Advance the streamed iterator and compute the new row's join key. + /** Advance the left iterator and compute the new row's join key. * * @return - * true if the streamed iterator returned a row and false otherwise. + * true if the left iterator returned a row and false otherwise. */ - private def advancedStreamed(): Boolean = { - if (streamedIter.advanceNext()) { - streamedRow = streamedIter.getRow - streamedRowEquiKey = streamedEquiKeyGenerator(streamedRow) - streamedRowPITKey = streamedPITKeyGenerator(streamedRow) + private def advancedLeft(): Boolean = { + if (leftIter.advanceNext()) { + leftRow = leftIter.getRow + leftRowEquiKey = leftEquiKeyGenerator(leftRow) + leftRowPITKey = leftPITKeyGenerator(leftRow) true } else { - streamedRow = null - streamedRowPITKey = null - streamedRowEquiKey = null + leftRow = null + leftRowPITKey = null + leftRowEquiKey = null false } } - /** Advance the buffered iterator until we find a row with join key + /** If returnNulls exit search early with `rightMatch = null`, otherwise + * advance the left iterator to continue searching. * * @return - * true if the buffered iterator returned a row and false otherwise. + * (keepSearching: Boolean, potentiallyMoreJoinRows: Boolean) */ - private def advancedBuffered(): Boolean = { - if (bufferedIter.advanceNext()) { - bufferedRow = bufferedIter.getRow - bufferedRowEquiKey = bufferedEquiKeyGenerator(bufferedRow) - bufferedRowPITKey = bufferedPITKeyGenerator(bufferedRow) - true - } else { - bufferedRow = null - bufferedRowEquiKey = null - bufferedRowPITKey = null + private def handleNoMatchForLeftRow(): (Boolean, Boolean) = { + if (returnNulls) { + rightMatch = null + return ( + false, // Stop searching to return the left row. + true // Potentially there will be another left row, therefore more join rows. + ) + } + return ( + advancedLeft(), // Stop searching if the left iterator is exhausted + // Only relevant if the search is stopped on this iteration. That will happen only + // if the steamed iterator is exhausted, so there will be no more join rows. false + ) + } + + /** Advance to the next right row. + * + * @return + * (keepSearching: Boolean, potentiallyMoreJoinRows: Boolean) + */ + private def handleRightRowDoesNotMatch(): (Boolean, Boolean) = { + rightMatch = null + return ( + advancedRightToRowWithNullFreeJoinKeys(), // Stop searching if right iterator is exhausted. + // If returnNulls there could be more unmatched join rows. + // If not returnNulls there can be no more join rows. + returnNulls + ) + } + + /** Advance the right iterator until we find a row with join keys + * + * @return + * true if the right iterator returned a row and false otherwise. + */ + private def advancedRightToRowWithNullFreeJoinKeys(): Boolean = { + var foundRow: Boolean = false + while (!foundRow && rightIter.advanceNext()) { + rightRow = rightIter.getRow + rightRowEquiKey = rightEquiKeyGenerator(rightRow) + rightRowPITKey = rightPITKeyGenerator(rightRow) + foundRow = !rightRowEquiKey.anyNull && !rightRowPITKey.anyNull + } + if (!foundRow) { + rightRow = null + rightRowEquiKey = null + rightRowPITKey = null + false + } else { + true } } } diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index b9271c1..816023f 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -24,11 +24,12 @@ package io.github.ackuq.pit +import org.apache.spark.sql.DataFrame import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.StructType import org.scalatest.flatspec.AnyFlatSpec -import org.apache.spark.sql.DataFrame + import EarlyStopSortMerge.pit import data.SmallDataSortMerge @@ -59,6 +60,9 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { tolerance: Int ): Unit = { + leftDataFrame.show() + rightDataFrame.show() + val pitJoin = leftDataFrame.join( rightDataFrame, @@ -70,6 +74,12 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { joinType ) + pitJoin.explain("codegen") + pitJoin.printSchema() + pitJoin.show() + + expectedDataFrame.show() + assert(pitJoin.schema.equals(expectedSchema)) assert(pitJoin.collect().sameElements(expectedDataFrame.collect())) } @@ -287,8 +297,8 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { ) { testSearchingBackwardForMatches( "inner", - smallData.fg3_with_nulls, - smallData.fg1_with_nulls, + smallData.fg3_with_value_nulls, + smallData.fg1_with_value_nulls, smallData.PIT_3_1_WITH_NULLS, smallData.PIT_2_NULLABLE_schema, 0 @@ -300,8 +310,8 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { ) { testSearchingBackwardForMatches( "left", - smallData.fg3_with_nulls, - smallData.fg1_with_nulls, + smallData.fg3_with_value_nulls, + smallData.fg1_with_value_nulls, smallData.PIT_3_1_WITH_NULLS_OUTER, smallData.PIT_2_NULLABLE_OUTER_schema, 0 @@ -313,8 +323,8 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { ) { testSearchingBackwardForMatches( "inner", - smallData.fg3_with_nulls, - smallData.fg1_with_nulls, + smallData.fg3_with_value_nulls, + smallData.fg1_with_value_nulls, smallData.PIT_3_1_T1_WITH_NULLS, smallData.PIT_2_NULLABLE_schema, 1 @@ -326,14 +336,69 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { ) { testSearchingBackwardForMatches( "left", - smallData.fg3_with_nulls, - smallData.fg1_with_nulls, + smallData.fg3_with_value_nulls, + smallData.fg1_with_value_nulls, smallData.PIT_3_1_T1_WITH_NULLS_OUTER, smallData.PIT_2_NULLABLE_OUTER_schema, 1 ) } + testBothCodegenAndInterpreted( + "inner_join_nulls_in_join_keys" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1_with_key_nulls, + smallData.fg3_with_key_nulls, + smallData.PIT_1_3_WITH_KEY_NULLS, + // It could be argued the correct schema would be `smallData.PIT_2_schema`, + // i.e. with join key columns made non-nullable. However, the normal spark inner + // join does not do this, so we don't either. + smallData.PIT_2_NULLABLE_KEYS_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_nulls_in_join_key" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1_with_key_nulls, + smallData.fg3_with_key_nulls, + smallData.PIT_1_3_WITH_KEY_NULLS_OUTER, + smallData.PIT_2_NULLABLE_KEYS_OUTER_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "inner_join_duplicate_join_keys" + ) { + testSearchingBackwardForMatches( + "inner", + smallData.fg1_duplicates, + smallData.fg3_duplicates, + smallData.PIT_1_3_DUPLICATES, + smallData.PIT_2_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_duplicate_join_keys" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg1_duplicates, + smallData.fg3_duplicates, + smallData.PIT_1_3_DUPLICATES, + smallData.PIT_2_OUTER_schema, + 0 + ) + } + def testJoiningThreeDataframes( joinType: String, expectedSchema: StructType diff --git a/scala/src/test/scala/data/SmallData.scala b/scala/src/test/scala/data/SmallData.scala index 5905c3b..3222012 100644 --- a/scala/src/test/scala/data/SmallData.scala +++ b/scala/src/test/scala/data/SmallData.scala @@ -41,13 +41,32 @@ class SmallData(spark: SparkSession) { Row(1, 7, "1y"), Row(2, 8, "2y") ) - val RAW_1_WITH_NULLS = Seq( + val RAW_1_DUPLICTES = Seq( + Row(1, 4, "1z"), + Row(1, 4, "1z"), + Row(1, 5, "1x"), + Row(1, 5, "1x"), + Row(2, 6, "2x"), + Row(2, 6, "2x"), + Row(1, 7, "1y"), + Row(1, 7, "1y"), + Row(2, 8, "2y"), + Row(2, 8, "2y") + ) + val RAW_1_WITH_VALUE_NULLS = Seq( Row(1, 4, "1z"), Row(1, 5, null), Row(2, 6, "2x"), Row(1, 7, "1y"), Row(2, 8, "2y") ) + val RAW_1_WITH_KEY_NULLS = Seq( + Row(1, 4, "1z"), + Row(1, null, "1x"), + Row(2, 6, "2x"), + Row(1, 7, "1y"), + Row(null, 8, "2y") + ) val RAW_3 = Seq( Row(1, 1, "f3-1-1"), Row(2, 2, "f3-2-2"), @@ -55,13 +74,32 @@ class SmallData(spark: SparkSession) { Row(2, 8, "f3-2-8"), Row(1, 10, "f3-1-10") ) - val RAW_3_WITH_NULLS = Seq( + val RAW_3_DUPLICATES = Seq( + Row(1, 1, "f3-1-1"), + Row(1, 1, "f3-1-1"), + Row(2, 2, "f3-2-2"), + Row(2, 2, "f3-2-2"), + Row(1, 6, "f3-1-6"), + Row(1, 6, "f3-1-6"), + Row(2, 8, "f3-2-8"), + Row(2, 8, "f3-2-8"), + Row(1, 10, "f3-1-10"), + Row(1, 10, "f3-1-10") + ) + val RAW_3_WITH_VALUE_NULLS = Seq( Row(1, 1, "f3-1-1"), Row(2, 2, "f3-2-2"), Row(1, 6, "f3-1-6"), Row(2, 8, null), Row(1, 10, "f3-1-10") ) + val RAW_3_WITH_KEY_NULLS = Seq( + Row(1, 1, "f3-1-1"), + Row(2, null, "f3-2-2"), + Row(null, 6, "f3-1-6"), + Row(2, 8, "f3-2-8"), + Row(1, 10, "f3-1-10") + ) val schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), @@ -69,7 +107,7 @@ class SmallData(spark: SparkSession) { StructField("value", StringType, nullable = false) ) ) - val schema_nullable: StructType = StructType( + val schema_value_nullable: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), @@ -77,21 +115,43 @@ class SmallData(spark: SparkSession) { ) ) + val schema_keys_nullable: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = false) + ) + ) + val fg1: DataFrame = spark.createDataFrame(spark.sparkContext.parallelize(RAW_1), schema) - val fg1_with_nulls: DataFrame = + val fg1_duplicates: DataFrame = + spark.createDataFrame(spark.sparkContext.parallelize(RAW_1_DUPLICTES), schema) + val fg1_with_value_nulls: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(RAW_1_WITH_NULLS), - schema_nullable + spark.sparkContext.parallelize(RAW_1_WITH_VALUE_NULLS), + schema_value_nullable + ) + val fg1_with_key_nulls: DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(RAW_1_WITH_KEY_NULLS), + schema_keys_nullable ) val fg2: DataFrame = spark.createDataFrame(spark.sparkContext.parallelize(RAW_1), schema) val fg3: DataFrame = spark.createDataFrame(spark.sparkContext.parallelize(RAW_3), schema) - val fg3_with_nulls: DataFrame = + val fg3_duplicates: DataFrame = + spark.createDataFrame(spark.sparkContext.parallelize(RAW_3_DUPLICATES), schema) + val fg3_with_value_nulls: DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(RAW_3_WITH_VALUE_NULLS), + schema_value_nullable + ) + val fg3_with_key_nulls: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(RAW_3_WITH_NULLS), - schema_nullable + spark.sparkContext.parallelize(RAW_3_WITH_KEY_NULLS), + schema_keys_nullable ) val empty: DataFrame = diff --git a/scala/src/test/scala/data/SmallDataSortMerge.scala b/scala/src/test/scala/data/SmallDataSortMerge.scala index aa2deb5..104ae4b 100644 --- a/scala/src/test/scala/data/SmallDataSortMerge.scala +++ b/scala/src/test/scala/data/SmallDataSortMerge.scala @@ -55,17 +55,40 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { Row(1, 5, "1x", 1, 1, "f3-1-1"), Row(1, 4, "1z", 1, 1, "f3-1-1") ) + val PIT_1_3_DUPLICATES_RAW = Seq( + Row(2, 8, "2y", 2, 8, "f3-2-8"), + Row(2, 8, "2y", 2, 8, "f3-2-8"), + Row(2, 6, "2x", 2, 2, "f3-2-2"), + Row(2, 6, "2x", 2, 2, "f3-2-2"), + Row(1, 7, "1y", 1, 6, "f3-1-6"), + Row(1, 7, "1y", 1, 6, "f3-1-6"), + Row(1, 5, "1x", 1, 1, "f3-1-1"), + Row(1, 5, "1x", 1, 1, "f3-1-1"), + Row(1, 4, "1z", 1, 1, "f3-1-1"), + Row(1, 4, "1z", 1, 1, "f3-1-1") + ) + val PIT_1_3_WITH_KEY_NULLS_RAW = Seq( + Row(1, 7, "1y", 1, 1, "f3-1-1"), + Row(1, 4, "1z", 1, 1, "f3-1-1") + ) + val PIT_1_3_WITH_KEY_NULLS_OUTER_RAW = Seq( + Row(2, 6, "2x", null, null, null), + Row(1, 7, "1y", 1, 1, "f3-1-1"), + Row(1, 4, "1z", 1, 1, "f3-1-1"), + Row(1, null, "1x", null, null, null), + Row(null, 8, "2y", null, null, null) + ) val PIT_3_1_RAW = Seq( Row(2, 8, "f3-2-8", 2, 8, "2y"), Row(1, 10, "f3-1-10", 1, 7, "1y"), Row(1, 6, "f3-1-6", 1, 5, "1x") ) - val PIT_3_1_WITH_NULLS_RAW = Seq( + val PIT_3_1_WITH_VALUE_NULLS_RAW = Seq( Row(2, 8, null, 2, 8, "2y"), Row(1, 10, "f3-1-10", 1, 7, "1y"), Row(1, 6, "f3-1-6", 1, 5, null) ) - val PIT_3_1_T1_WITH_NULLS_RAW = Seq( + val PIT_3_1_T1_WITH_VALUE_NULLS_RAW = Seq( Row(2, 8, null, 2, 8, "2y"), Row(1, 6, "f3-1-6", 1, 5, null) ) @@ -76,14 +99,14 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { Row(1, 6, "f3-1-6", 1, 5, "1x"), Row(1, 1, "f3-1-1", null, null, null) ) - val PIT_3_1_T1_WITH_NULLS_OUTER_RAW = Seq( + val PIT_3_1_T1_WITH_VALUE_NULLS_OUTER_RAW = Seq( Row(2, 8, null, 2, 8, "2y"), Row(2, 2, "f3-2-2", null, null, null), Row(1, 10, "f3-1-10", null, null, null), Row(1, 6, "f3-1-6", 1, 5, null), Row(1, 1, "f3-1-1", null, null, null) ) - val PIT_3_1_WITH_NULLS_OUTER_RAW = Seq( + val PIT_3_1_WITH_VALUE_NULLS_OUTER_RAW = Seq( Row(2, 8, null, 2, 8, "2y"), Row(2, 2, "f3-2-2", null, null, null), Row(1, 10, "f3-1-10", 1, 7, "1y"), @@ -128,6 +151,18 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { StructField("value", StringType, nullable = true) ) ) + + val PIT_2_NULLABLE_KEYS_schema: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = false), + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = false) + ) + ) + val PIT_2_OUTER_schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), @@ -149,6 +184,17 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { ) ) + val PIT_2_NULLABLE_KEYS_OUTER_schema: StructType = StructType( + Seq( + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = false), + StructField("id", IntegerType, nullable = true), + StructField("ts", IntegerType, nullable = true), + StructField("value", StringType, nullable = true) + ) + ) + val PIT_3_schema: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), @@ -183,6 +229,20 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { spark.sparkContext.parallelize(PIT_1_3_RAW), PIT_2_schema ) + val PIT_1_3_DUPLICATES: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_1_3_DUPLICATES_RAW), + PIT_2_schema + ) + + val PIT_1_3_WITH_KEY_NULLS: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_1_3_WITH_KEY_NULLS_RAW), + PIT_2_schema + ) + + val PIT_1_3_WITH_KEY_NULLS_OUTER: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_1_3_WITH_KEY_NULLS_OUTER_RAW), + PIT_2_NULLABLE_KEYS_OUTER_schema + ) val PIT_1_3_T1: DataFrame = spark.createDataFrame( spark.sparkContext.parallelize(PIT_1_3_T1_RAW), @@ -205,22 +265,22 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { ) val PIT_3_1_WITH_NULLS: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(PIT_3_1_WITH_NULLS_RAW), + spark.sparkContext.parallelize(PIT_3_1_WITH_VALUE_NULLS_RAW), PIT_2_NULLABLE_schema ) val PIT_3_1_T1_WITH_NULLS: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(PIT_3_1_T1_WITH_NULLS_RAW), + spark.sparkContext.parallelize(PIT_3_1_T1_WITH_VALUE_NULLS_RAW), PIT_2_NULLABLE_schema ) val PIT_3_1_WITH_NULLS_OUTER: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(PIT_3_1_WITH_NULLS_OUTER_RAW), + spark.sparkContext.parallelize(PIT_3_1_WITH_VALUE_NULLS_OUTER_RAW), PIT_2_NULLABLE_OUTER_schema ) val PIT_3_1_T1_WITH_NULLS_OUTER: DataFrame = spark.createDataFrame( - spark.sparkContext.parallelize(PIT_3_1_T1_WITH_NULLS_OUTER_RAW), + spark.sparkContext.parallelize(PIT_3_1_T1_WITH_VALUE_NULLS_OUTER_RAW), PIT_2_NULLABLE_OUTER_schema ) From e29716977c049755f87a78f904009f39f395f123 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 15 Jan 2024 13:55:07 +0000 Subject: [PATCH 05/19] Update README to reflect new way of initialising the context (#7) --- README.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index e2ccbe9..54ab795 100644 --- a/README.md +++ b/README.md @@ -62,11 +62,10 @@ pip install spark-pit The object `PitContext` is the entrypoint for all of the functionality of the lirary. You can initialize this context with the following code: ```py -from pyspark import SQLContext from ackuq.pit import PitContext -sql_context = SQLContext(spark.sparkContext) -pit_context = PitContext(sql_context) +spark = +pit_context = PitContext(spark) ``` ### 2. Performing a PIT join From 8c9256fee4951c7203de3e97000e222e2b20321e Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 15 Jan 2024 13:58:58 +0000 Subject: [PATCH 06/19] Update python spark to 3.5.0 (#8) --- python/requirements_common.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/requirements_common.txt b/python/requirements_common.txt index 3f8a5ac..67e0899 100644 --- a/python/requirements_common.txt +++ b/python/requirements_common.txt @@ -1,3 +1,3 @@ # Common requirements used for all environments -pyspark==3.2.1 \ No newline at end of file +pyspark==3.5.0 From 832dbe8873264c6c50483eec7cf8262bc2277c66 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 15 Jan 2024 17:57:59 +0000 Subject: [PATCH 07/19] Bump versions (#9) --- python/VERSION | 2 +- scala/VERSION | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/VERSION b/python/VERSION index 79a2734..ab8684e 100644 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.5.0 \ No newline at end of file +0.6.0-tom diff --git a/scala/VERSION b/scala/VERSION index 5d4294b..ab8684e 100644 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.5.1 \ No newline at end of file +0.6.0-tom From 67a454ee9426e517d58becfb4efb65b6232ffbc3 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Wed, 17 Jan 2024 11:49:37 +0000 Subject: [PATCH 08/19] Fully remove non equi non pit condition (#10) * Remove `condition` from exec code * Fail physical planning if there are non-equi conditions * Add test * Throw an exception --- .../main/scala/execution/PITJoinExec.scala | 61 +------------------ scala/src/main/scala/execution/Patterns.scala | 13 ++-- .../src/test/scala/EarlyStopMergeTests.scala | 17 +++++- 3 files changed, 22 insertions(+), 69 deletions(-) diff --git a/scala/src/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index fc750f3..872e6cb 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -553,32 +553,6 @@ protected[pit] case class PITJoinExec( (leftRow, matched) } - /** Splits variables based on whether it's used by condition or not, returns - * the code to create these variables before the condition and after the - * condition. - * - * Only a few columns are used by condition, then we can skip the accessing - * of those columns that are not used by condition also filtered out by - * condition. - */ - private def splitVarsByCondition( - attributes: Seq[Attribute], - variables: Seq[ExprCode] - ): (String, String) = { - if (condition.isDefined) { - val condRefs = condition.get.references - val (used, notUsed) = - attributes.zip(variables).partition { case (a, ev) => - condRefs.contains(a) - } - val beforeCond = evaluateVariables(used.map(_._2)) - val afterCond = evaluateVariables(notUsed.map(_._2)) - (beforeCond, afterCond) - } else { - (evaluateVariables(variables), "") - } - } - override def needCopyResult: Boolean = true override protected def doProduce(ctx: CodegenContext): String = { @@ -605,38 +579,7 @@ protected[pit] case class PITJoinExec( val numOutput = metricTerm(ctx, "numOutputRows") - val (beforeLoop, condCheck) = if (condition.isDefined) { - // Split the code of creating variables based on whether it's used by condition or not. - val loaded = ctx.freshName("loaded") - val (leftBefore, leftAfter) = splitVarsByCondition(left.output, leftVars) - val (rightBefore, rightAfter) = - splitVarsByCondition(right.output, rightVars) - // Generate code for condition - ctx.currentVars = leftVars ++ rightVars - val cond = - BindReferences.bindReference(condition.get, output).genCode(ctx) - // evaluate the columns those used by condition before loop - val before = - s""" - |boolean $loaded = false; - |$leftBefore - """.stripMargin - - val checking = - s""" - |$rightBefore - |${cond.code} - |if (${cond.isNull}|| !${cond.value}) continue; - |if (!$loaded) { - | $loaded = true; - | $leftAfter - |} - |$rightAfter - """.stripMargin - (before, checking) - } else { - (evaluateVariables(leftVars), "") - } + val beforeLoop = evaluateVariables(leftVars) val thisPlan = ctx.addReferenceObj("plan", this) val eagerCleanup = s"$thisPlan.cleanupResources();" @@ -647,7 +590,6 @@ protected[pit] case class PITJoinExec( | ${leftVarDecl.mkString("\n")} | ${beforeLoop.trim} | InternalRow $rightRow = (InternalRow) $matched; - | ${condCheck.trim} | $numOutput.add(1); | ${consume(ctx, leftVars ++ rightVars)}; | if (shouldStop()) return; @@ -659,7 +601,6 @@ protected[pit] case class PITJoinExec( | ${leftVarDecl.mkString("\n")} | ${beforeLoop.trim} | InternalRow $rightRow = (InternalRow) $matched; - | ${condCheck.trim} | $numOutput.add(1); | ${consume(ctx, leftVars ++ rightVars)} | if (shouldStop()) return; diff --git a/scala/src/main/scala/execution/Patterns.scala b/scala/src/main/scala/execution/Patterns.scala index 15cb092..9e62d2e 100644 --- a/scala/src/main/scala/execution/Patterns.scala +++ b/scala/src/main/scala/execution/Patterns.scala @@ -108,13 +108,10 @@ object PITJoinExtractEquality extends ExtractEqualityKeys { // These need to be sortable in order to make the algorithm work optimized val equiJoinKeys = getEquiJoinKeys(predicates, join.left, join.right) - val otherPredicates = predicates.filterNot { - case EqualTo(l, r) if l.references.isEmpty || r.references.isEmpty => - false - case Equality(l, r) => - canEvaluate(l, join.left) && canEvaluate(r, join.right) || - canEvaluate(l, join.right) && canEvaluate(r, join.left) - case _ => false + if (predicates.length != equiJoinKeys.length) { + throw new IllegalArgumentException( + "Besides the PIT key, only equi-conditions are supported for PIT joins" + ) } val leftPitKey = if (canEvaluate(join.pitCondition.children.head, join.left)) @@ -135,7 +132,7 @@ object PITJoinExtractEquality extends ExtractEqualityKeys { rightPitKey, leftEquiKeys, rightEquiKeys, - otherPredicates.reduceOption(And), + None, join.returnNulls, join.tolerance, join.left, diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index 816023f..47fbc56 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -352,7 +352,7 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { smallData.fg1_with_key_nulls, smallData.fg3_with_key_nulls, smallData.PIT_1_3_WITH_KEY_NULLS, - // It could be argued the correct schema would be `smallData.PIT_2_schema`, + // It could be argued the correct schema would be `smallData.PIT_2_schema`, // i.e. with join key columns made non-nullable. However, the normal spark inner // join does not do this, so we don't either. smallData.PIT_2_NULLABLE_KEYS_schema, @@ -431,4 +431,19 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { testBothCodegenAndInterpreted("left_join_three_dataframes") { testJoiningThreeDataframes("left", smallData.PIT_3_OUTER_schema) } + + testBothCodegenAndInterpreted("fail_during_planning_for_non_equi_condition") { + val fg1 = smallData.fg1 + val fg2 = smallData.fg2 + + val pitJoin = + fg1.join( + fg2, + pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") && fg1("value") > fg2("value"), + "inner" + ) + intercept[IllegalArgumentException] { + pitJoin.explain() + } + } } From 40ef9624f549b8ef7476e11b3f1331adb9511f50 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Thu, 18 Jan 2024 08:49:42 +0000 Subject: [PATCH 09/19] Fix invalid query pushdown (#11) * Initial testcase * Correct test * Attempt to create a new version of normal `.join` * Delete mess * Switch to new spark extensions framework * First explicit pitJoin compiles * In progress: create new entrypoint to PIT join * Working test * All scala tests working * Create python side joinPIT to replace context * Rename some paramaters * Update parameter order in tests * Orgnise imports * Tidy --- python/ackuq/pit/__init__.py | 2 +- python/ackuq/pit/context.py | 187 ------------------ python/ackuq/pit/joinPIT.py | 46 +++++ python/tests/test_sort_merge.py | 43 ++-- python/tests/utils.py | 23 +-- scala/src/main/scala/EarlyStopSortMerge.scala | 108 +++++++--- scala/src/main/scala/execution/Patterns.scala | 17 +- scala/src/main/scala/logical/PITJoin.scala | 3 +- scala/src/main/scala/logical/PITRule.scala | 100 ---------- .../src/test/scala/EarlyStopMergeTests.scala | 92 ++++++--- .../test/scala/SparkSessionTestWrapper.scala | 1 + .../test/scala/data/SmallDataSortMerge.scala | 8 + 12 files changed, 235 insertions(+), 395 deletions(-) delete mode 100644 python/ackuq/pit/context.py create mode 100644 python/ackuq/pit/joinPIT.py delete mode 100644 scala/src/main/scala/logical/PITRule.scala diff --git a/python/ackuq/pit/__init__.py b/python/ackuq/pit/__init__.py index fccc5be..ebb0f55 100644 --- a/python/ackuq/pit/__init__.py +++ b/python/ackuq/pit/__init__.py @@ -22,4 +22,4 @@ # SOFTWARE. # -from ackuq.pit.context import PitContext # noqa: F401 +from ackuq.pit.joinPIT import joinPIT # noqa: F401 diff --git a/python/ackuq/pit/context.py b/python/ackuq/pit/context.py deleted file mode 100644 index 862c916..0000000 --- a/python/ackuq/pit/context.py +++ /dev/null @@ -1,187 +0,0 @@ -# -# MIT License -# -# Copyright (c) 2022 Axel Pettersson -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# - -from typing import Any, List, Optional, Tuple - -from py4j.java_gateway import JavaPackage -from pyspark.sql import Column, DataFrame, SparkSession -from pyspark.sql.column import _to_java_column, _to_seq # type: ignore -from pyspark.sql.functions import lit - - -class PitContext(object): - """ - Entrypoint for all the functionality. - Essentially a wrapper for the Scala functions. - """ - - def _init_early_stop_sort_merge(self) -> None: - # Call the init function in the Scala object, - # will register the strategies and rules - self._essm.init(self._spark._jsparkSession) # type: ignore - - CLASSPATH_ERROR_MSG = ( - "Java class {} could not be imported, check that it is included in the JVM" - " classpaths." - ) - - def _check_classpath(self): - # If the classpath does not exist py4j will assume - # that the references are Java packages instead - - if isinstance(self._essm, JavaPackage): - raise ImportError( - self.CLASSPATH_ERROR_MSG.format( - "io.github.ackuq.pit.EarlyStopSortMerge" - ) - ) - - if isinstance(self._union, JavaPackage): - raise ImportError( - self.CLASSPATH_ERROR_MSG.format("io.github.ackuq.pit.Union") - ) - - if isinstance(self._exploding, JavaPackage): - raise ImportError( - self.CLASSPATH_ERROR_MSG.format("io.github.ackuq.pit.Exploding") - ) - - def __init__(self, spark: SparkSession) -> None: - self._spark = spark - self._sc = self._spark.sparkContext - self._jsc = self._sc._jsc - self._jvm = self._sc._jvm - - self._essm = self._jvm.io.github.ackuq.pit.EarlyStopSortMerge - self._union = self._jvm.io.github.ackuq.pit.Union - self._exploding = self._jvm.io.github.ackuq.pit.Exploding - self._check_classpath() - self._init_early_stop_sort_merge() - - def _to_scala_seq(self, list_like: List): - """ - Converts a list to a Scala sequence. - """ - return ( - self._jvm.scala.collection.JavaConverters.asScalaBufferConverter(list_like) - .asScala() - .toSeq() - ) - - def _to_scala_tuple(self, tuple: Tuple): - """ - Converts a Python tuple to Scala tuple - """ - return self._jvm.__getattr__("scala.Tuple" + str(len(tuple)))(*tuple) - - def _scala_none(self): - return self._jvm.scala.__getattr__("None$").__getattr__("MODULE$") - - def _to_scala_option(self, x: Optional[Any]): - return self._jvm.scala.Some(x) if x is not None else self._scala_none() - - def pit_udf(self, left: Column, right: Column, tolerance: int = 0): - """ - Used for executing an early stop sort merge join by using custom UDF. - - :param left The timestamp column of left table - :param right The timestamp column of right table - :param tolerance Optional, The tolerance of the PIT result - """ - _pit_udf = self._essm.getPit() - return Column( - _pit_udf.apply( - _to_seq( - self._sc, - [left, right, lit(tolerance)], - _to_java_column, - ) - ) - ) - - def union( - self, - left: DataFrame, - right: DataFrame, - right_prefix: str, - left_ts_column: str = "ts", - right_ts_column: str = "ts", - left_prefix: Optional[str] = None, - partition_cols: List[str] = [], - ): - """ - Perform a backward asof join using the left table for event times. - - :param left The left dataframe, will be used as reference - :param right The right dataframe, will be used to merge - :param right_prefix The prefix for the right columns in result - :param left_ts_column The column used for timestamps in left DF - :param right_ts_column The column used for timestamps in right DF - :param left_prefix Optional, the prefix for the left columns in result - :param partition_cols The columns used for partitioning, if used - """ - - return DataFrame( - self._union.join( - left._jdf, - right._jdf, - left_ts_column, - right_ts_column, - self._to_scala_option(left_prefix), - right_prefix, - self._to_scala_seq(partition_cols), - ), - self._sql_context, - ) - - def exploding( - self, - left: DataFrame, - right: DataFrame, - left_ts_column: Column, - right_ts_column: Column, - partition_cols: List[Tuple[Column, Column]] = [], - ): - """ - Perform a backward asof join using the left table for event times. - - :param left The left dataframe, will be used as reference - :param right The right dataframe, will be used to merge - :param left_ts_column The column used for timestamps in left DF - :param right_ts_column The column used for timestamps in right DF - :param partition_cols The columns used for partitioning, if used - """ - _partition_cols = map( - lambda p: self._to_scala_tuple((p[0]._jc, p[1]._jc)), partition_cols - ) - return DataFrame( - self._exploding.join( - left._jdf, - right._jdf, - left_ts_column._jc, # type: ignore - right_ts_column._jc, # type: ignore - self._to_scala_seq(list(_partition_cols)), - ), - self._sql_context, - ) diff --git a/python/ackuq/pit/joinPIT.py b/python/ackuq/pit/joinPIT.py new file mode 100644 index 0000000..43fd830 --- /dev/null +++ b/python/ackuq/pit/joinPIT.py @@ -0,0 +1,46 @@ +# +# MIT License +# +# Copyright (c) 2022 Axel Pettersson +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# + +from pyspark.sql import Column, DataFrame + + +def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitKey: Column, + rightPitKey: Column, + on: Column, + how: str = "inner", + tolerance: int = 0, +) -> DataFrame: + jdf = left.sparkSession.sparkContext._jvm.io.github.ackuq.pit.EarlyStopSortMerge.joinPIT( + left._jdf, + right._jdf, + leftPitKey._jc, + rightPitKey._jc, + on._jc, + how, + tolerance, + ) + return DataFrame(jdf, left.sparkSession) diff --git a/python/tests/test_sort_merge.py b/python/tests/test_sort_merge.py index df29a09..7c4fe59 100644 --- a/python/tests/test_sort_merge.py +++ b/python/tests/test_sort_merge.py @@ -22,23 +22,22 @@ # SOFTWARE. # +from ackuq.pit.joinPIT import joinPIT from tests.data import SmallDataSortMerge from tests.utils import SparkTests class SortMergeUnionAsOfTest(SparkTests): - def setUp(self) -> None: - super().setUp() - self.small_data = SmallDataSortMerge(self.spark) + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + cls.small_data = SmallDataSortMerge(cls.spark) def test_two_aligned(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg2 - pit_join = fg1.join( - fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"]) & (fg1["id"] == fg2["id"]), - ) + pit_join = joinPIT(fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"])) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_2.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_2.collect()) @@ -47,10 +46,7 @@ def test_two_misaligned(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = fg1.join( - fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"]) & (fg1["id"] == fg2["id"]), - ) + pit_join = joinPIT(fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"])) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3.collect()) @@ -60,15 +56,15 @@ def test_three_misaligned(self): fg2 = self.small_data.fg2 fg3 = self.small_data.fg3 - left = fg1.join( + left = joinPIT( + fg1, fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"]) & (fg1["id"] == fg2["id"]), + fg1["ts"], + fg2["ts"], + (fg1["id"] == fg2["id"]), ) - pit_join = left.join( - fg3, - self.pit_context.pit_udf(fg1["ts"], fg3["ts"]) & (fg1["id"] == fg3["id"]), - ) + pit_join = joinPIT(left, fg3, fg1["ts"], fg3["ts"], (fg1["id"] == fg3["id"])) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_2_3.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_2_3.collect()) @@ -77,10 +73,8 @@ def test_two_tolerance(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = fg1.join( - fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"], 1) - & (fg1["id"] == fg2["id"]), + pit_join = joinPIT( + fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]), tolerance=1 ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3_T1.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3_T1.collect()) @@ -89,11 +83,8 @@ def test_two_tolerance_outer(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = fg1.join( - fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"], 1) - & (fg1["id"] == fg2["id"]), - "left", + pit_join = joinPIT( + fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]), "left", 1 ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3_T1_OUTER.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3_T1_OUTER.collect()) diff --git a/python/tests/utils.py b/python/tests/utils.py index 5b6b3bb..8d7d1d8 100644 --- a/python/tests/utils.py +++ b/python/tests/utils.py @@ -25,29 +25,30 @@ import os import unittest -from pyspark import SQLContext from pyspark.sql import SparkSession from pyspark.sql.types import StructField, StructType -from ackuq.pit.context import PitContext - class SparkTests(unittest.TestCase): - def setUp(self) -> None: - self.jar_location = os.environ["SCALA_PIT_JAR"] - print("Loading jar from location: {}".format(self.jar_location)) - self.spark = ( + spark: SparkSession + + @classmethod + def setUpClass(cls) -> None: + jar_location = os.environ["SCALA_PIT_JAR"] + print("Loading jar from location: {}".format(jar_location)) + cls.spark = ( SparkSession.builder.appName("sparkTests") .master("local") .config("spark.ui.showConsoleProgress", False) - .config("spark.driver.extraClassPath", self.jar_location) + .config("spark.driver.extraClassPath", jar_location) .config("spark.sql.shuffle.partitions", 1) + .config("spark.sql.extensions", "io.github.ackuq.pit.SparkPIT") .getOrCreate() ) - self.pit_context = PitContext(self.spark) - def tearDown(self) -> None: - self.spark.stop() + @classmethod + def tearDownClass(cls) -> None: + cls.spark.stop() def _assertFieldsEqual(self, a: StructField, b: StructField): self.assertEqual(a.name.lower(), b.name.lower()) diff --git a/scala/src/main/scala/EarlyStopSortMerge.scala b/scala/src/main/scala/EarlyStopSortMerge.scala index 6853bfd..c20dd03 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -24,36 +24,92 @@ package io.github.ackuq.pit -import execution.CustomStrategy -import logical.PITRule +import org.apache.spark.sql.{ + Column, + DataFrame, + SparkSessionExtensionsProvider, + SparkSessionExtensions +} +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, JoinType} -import org.apache.spark.sql.expressions.UserDefinedFunction -import org.apache.spark.sql.functions.udf -import org.apache.spark.sql.{Column, SparkSession} +import execution.CustomStrategy +import logical.PITJoin object EarlyStopSortMerge { - private final val PIT_FUNCTION = ( - _: Column, - _: Column, - _: Long - ) => true - final val PIT_UDF_NAME = "PIT" - final val pit: UserDefinedFunction = udf(PIT_FUNCTION).withName(PIT_UDF_NAME) - - def init(spark: SparkSession): Unit = { - if (!spark.catalog.functionExists(PIT_UDF_NAME)) { - spark.udf.register(PIT_UDF_NAME, pit) - } - if (!spark.experimental.extraStrategies.contains(CustomStrategy)) { - spark.experimental.extraStrategies = - spark.experimental.extraStrategies :+ CustomStrategy - } - if (!spark.experimental.extraOptimizations.contains(PITRule)) { - spark.experimental.extraOptimizations = - spark.experimental.extraOptimizations :+ PITRule + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitExpression: Column, + rightPitExpression: Column, + joinType: String, + tolerance: Long + ): DataFrame = joinPIT( + left, + right, + leftPitExpression, + rightPitExpression, + None, + joinType, + tolerance + ) + + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitExpression: Column, + rightPitExpression: Column, + joinExprs: Column, + joinType: String, + tolerance: Long + ): DataFrame = joinPIT( + left, + right, + leftPitExpression, + rightPitExpression, + Some(joinExprs), + joinType, + tolerance + ) + + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitExpression: Column, + rightPitExpression: Column, + joinExprs: Option[Column], + joinType: String, + tolerance: Long + ): DataFrame = { + + val parsedJoinType = JoinType(joinType) + parsedJoinType match { + case LeftOuter | Inner => () + case x => + throw new IllegalArgumentException( + s"Join type $x not supported for PIT joins" + ) } + + val logicalPlan = PITJoin( + left.queryExecution.analyzed, + right.queryExecution.analyzed, + leftPitExpression.expr, + rightPitExpression.expr, + parsedJoinType == LeftOuter, + tolerance, + joinExprs.map(_.expr) + ) + new DataFrame( + left.sparkSession, + logicalPlan, + ExpressionEncoder(logicalPlan.schema) + ) } +} - // For the PySpark API - def getPit: UserDefinedFunction = pit +class SparkPIT extends SparkSessionExtensionsProvider { + override def apply(extensions: SparkSessionExtensions): Unit = { + extensions.injectPlannerStrategy(session => CustomStrategy) + } } diff --git a/scala/src/main/scala/execution/Patterns.scala b/scala/src/main/scala/execution/Patterns.scala index 9e62d2e..7004257 100644 --- a/scala/src/main/scala/execution/Patterns.scala +++ b/scala/src/main/scala/execution/Patterns.scala @@ -99,7 +99,7 @@ object PITJoinExtractEquality extends ExtractEqualityKeys { ) def unapply(join: PITJoin): Option[ReturnType] = { - logDebug(s"Considering join on: ${join.condition}") + logDebug(s"Considering PIT join") val predicates = join.condition.map(splitConjunctivePredicates).getOrElse(Nil) @@ -113,23 +113,14 @@ object PITJoinExtractEquality extends ExtractEqualityKeys { "Besides the PIT key, only equi-conditions are supported for PIT joins" ) } - val leftPitKey = - if (canEvaluate(join.pitCondition.children.head, join.left)) - join.pitCondition.children.head - else join.pitCondition.children(1) - - val rightPitKey = - if (canEvaluate(join.pitCondition.children.head, join.right)) - join.pitCondition.children.head - else join.pitCondition.children(1) val (leftEquiKeys, rightEquiKeys) = equiJoinKeys.unzip - logDebug(s"leftPitKey:$leftPitKey | rightPitKey:$rightPitKey") + logDebug(s"leftPitKey:${join.leftPitKey} | rightPitKey:${join.rightPitKey}") Some( ( - leftPitKey, - rightPitKey, + join.leftPitKey, + join.rightPitKey, leftEquiKeys, rightEquiKeys, None, diff --git a/scala/src/main/scala/logical/PITJoin.scala b/scala/src/main/scala/logical/PITJoin.scala index f4e37d8..93189e2 100644 --- a/scala/src/main/scala/logical/PITJoin.scala +++ b/scala/src/main/scala/logical/PITJoin.scala @@ -45,7 +45,8 @@ protected[pit] case object PITJoinType extends CustomJoinType { protected[pit] case class PITJoin( left: LogicalPlan, right: LogicalPlan, - pitCondition: Expression, + leftPitKey: Expression, + rightPitKey: Expression, returnNulls: Boolean, tolerance: Long, condition: Option[Expression] diff --git a/scala/src/main/scala/logical/PITRule.scala b/scala/src/main/scala/logical/PITRule.scala deleted file mode 100644 index d199c4a..0000000 --- a/scala/src/main/scala/logical/PITRule.scala +++ /dev/null @@ -1,100 +0,0 @@ -/* - * MIT License - * - * Copyright (c) 2022 Axel Pettersson - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -package io.github.ackuq.pit -package logical - -import EarlyStopSortMerge.PIT_UDF_NAME - -import org.apache.spark.sql.catalyst.expressions.{ - And, - Expression, - GreaterThanOrEqual, - PredicateHelper, - ScalaUDF -} -import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter} -import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan} -import org.apache.spark.sql.catalyst.rules.Rule - -protected[pit] object PITRule extends Rule[LogicalPlan] with PredicateHelper { - def apply(logicalPlan: LogicalPlan): LogicalPlan = - logicalPlan.transform { - case j @ Join(left, right, joinType, condition, _) => - val predicates = { - condition.map(splitConjunctivePredicates).getOrElse(Nil) - } - - val pitExpressions = getPITExpression(predicates) - if (pitExpressions.nonEmpty && pitExpressions.length == 1) { - joinType match { - case LeftOuter | Inner => () - case x => - throw new IllegalArgumentException( - s"Join type $x not supported for PIT joins" - ) - } - - val pitExpression: ScalaUDF = pitExpressions.head - val otherPredicates = predicates.filterNot { - case f: ScalaUDF if f.udfName.getOrElse("") == PIT_UDF_NAME => true - case _ => false - } - val pitCondition = - GreaterThanOrEqual( - pitExpression.children.head, - pitExpression.children(1) - ) - - val tolerance = - pitExpression - .children(2) - .eval() - .asInstanceOf[Long] - - // When join type is left outer, return nulls as well - val returnNulls = joinType == LeftOuter - - PITJoin( - left, - right, - pitCondition, - returnNulls, - tolerance, - otherPredicates.reduceOption(And) - ) - } else { - j - } - } - - private def getPITExpression(predicates: Seq[Expression]) = { - predicates.flatMap { - // TODO: Identifying by name is unsafe, maybe could be improved - case f: ScalaUDF if f.udfName.getOrElse("") == PIT_UDF_NAME => - Some(f) - case _ => None - } - } -} diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index 47fbc56..ecbc0ca 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -30,11 +30,11 @@ import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.StructType import org.scalatest.flatspec.AnyFlatSpec -import EarlyStopSortMerge.pit +import EarlyStopSortMerge.joinPIT import data.SmallDataSortMerge +import io.github.ackuq.pit.execution.CustomStrategy class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { - EarlyStopSortMerge.init(spark) val smallData = new SmallDataSortMerge(spark) def testBothCodegenAndInterpreted(name: String)(f: => Unit): Unit = { @@ -63,16 +63,15 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { leftDataFrame.show() rightDataFrame.show() - val pitJoin = - leftDataFrame.join( - rightDataFrame, - pit( - leftDataFrame("ts"), - rightDataFrame("ts"), - lit(tolerance) - ) && leftDataFrame("id") === rightDataFrame("id"), - joinType - ) + val pitJoin = joinPIT( + leftDataFrame, + rightDataFrame, + leftDataFrame("ts"), + rightDataFrame("ts"), + leftDataFrame("id") === rightDataFrame("id"), + joinType, + tolerance, + ) pitJoin.explain("codegen") pitJoin.printSchema() @@ -399,6 +398,30 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { ) } + def testJoinThenFilter() { + val leftDataFrame = smallData.fg3 + val rightDataFrame = smallData.fg1 + + var pitJoin = joinPIT( + leftDataFrame, + rightDataFrame, + leftDataFrame("ts"), + rightDataFrame("ts"), + leftDataFrame("id") === rightDataFrame("id"), + "inner", + 0, + ).filter(rightDataFrame("value") === lit("1x")) + + assert(pitJoin.schema.equals(smallData.PIT_2_schema)) + assert(pitJoin.collect().sameElements(smallData.PIT_3_1_FILTERED.collect())) + } + + testBothCodegenAndInterpreted( + "inner_join_then_filter" + ) { + testJoinThenFilter() + } + def testJoiningThreeDataframes( joinType: String, expectedSchema: StructType @@ -407,19 +430,25 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { val fg2 = smallData.fg2 val fg3 = smallData.fg3 - val left = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id"), - joinType - ) + val left = joinPIT( + fg1, + fg2, + fg1("ts"), + fg2("ts"), + fg1("id") === fg2("id"), + joinType, + 0, + ) - val pitJoin = - left.join( - fg3, - pit(fg1("ts"), fg3("ts"), lit(0)) && fg1("id") === fg3("id"), - joinType - ) + val pitJoin = joinPIT( + left, + fg3, + fg1("ts"), + fg3("ts"), + fg1("id") === fg3("id"), + joinType, + 0, + ) assert(pitJoin.schema.equals(expectedSchema)) assert(pitJoin.collect().sameElements(smallData.PIT_1_2_3.collect())) @@ -436,12 +465,15 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { val fg1 = smallData.fg1 val fg2 = smallData.fg2 - val pitJoin = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") && fg1("value") > fg2("value"), - "inner" - ) + val pitJoin = joinPIT( + fg1, + fg2, + fg1("ts"), + fg2("ts"), + (fg1("id") === fg2("id") && fg1("value") > fg2("value")), + "inner", + 0, + ) intercept[IllegalArgumentException] { pitJoin.explain() } diff --git a/scala/src/test/scala/SparkSessionTestWrapper.scala b/scala/src/test/scala/SparkSessionTestWrapper.scala index 51bbcfd..c6b6d6b 100644 --- a/scala/src/test/scala/SparkSessionTestWrapper.scala +++ b/scala/src/test/scala/SparkSessionTestWrapper.scala @@ -33,6 +33,7 @@ trait SparkSessionTestWrapper { .appName("Spark PIT Tests") .config("spark.ui.showConsoleProgress", value = false) .config("spark.sql.shuffle.partitions", 1) + .config("spark.sql.extensions", "io.github.ackuq.pit.SparkPIT") .getOrCreate() spark.sparkContext.setLogLevel("ERROR") diff --git a/scala/src/test/scala/data/SmallDataSortMerge.scala b/scala/src/test/scala/data/SmallDataSortMerge.scala index 104ae4b..2eeed23 100644 --- a/scala/src/test/scala/data/SmallDataSortMerge.scala +++ b/scala/src/test/scala/data/SmallDataSortMerge.scala @@ -83,6 +83,9 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { Row(1, 10, "f3-1-10", 1, 7, "1y"), Row(1, 6, "f3-1-6", 1, 5, "1x") ) + val PIT_3_1_FILTERED_RAW = Seq( + Row(1, 6, "f3-1-6", 1, 5, "1x") + ) val PIT_3_1_WITH_VALUE_NULLS_RAW = Seq( Row(2, 8, null, 2, 8, "2y"), Row(1, 10, "f3-1-10", 1, 7, "1y"), @@ -264,6 +267,11 @@ class SmallDataSortMerge(spark: SparkSession) extends SmallData(spark) { PIT_2_OUTER_schema ) + val PIT_3_1_FILTERED: DataFrame = spark.createDataFrame( + spark.sparkContext.parallelize(PIT_3_1_FILTERED_RAW), + PIT_2_schema + ) + val PIT_3_1_WITH_NULLS: DataFrame = spark.createDataFrame( spark.sparkContext.parallelize(PIT_3_1_WITH_VALUE_NULLS_RAW), PIT_2_NULLABLE_schema From d098f144d4befbb8a454814b40700ffe42d5e302 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Thu, 18 Jan 2024 13:16:58 +0000 Subject: [PATCH 10/19] Update Readme (#12) * Make spark session an explicit argument * Update Readme --- README.md | 101 ++++---------------------------- python/README.md | 70 ---------------------- python/ackuq/pit/joinPIT.py | 7 ++- python/tests/test_sort_merge.py | 30 ++++++++-- 4 files changed, 41 insertions(+), 167 deletions(-) diff --git a/README.md b/README.md index 54ab795..aafa5ac 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,6 @@ Apart from utilising existing high-level implementations, a couple of implementa The thesis that this project laid the foundation for can be found here: http://www.diva-portal.org/smash/get/diva2:1695672/FULLTEXT01.pdf. -# This project is not being actively developed or maintained. If you wish to continue development on this project, feel free to make a fork! ## Table of contents @@ -15,11 +14,7 @@ The thesis that this project laid the foundation for can be found here: http://w - [Scala (Spark)](#scala-spark) - [Python (PySpark)](#python-pyspark) - [Quickstart (Python)](#quickstart-python) - 1. [Creating the context](#1-creating-the-context) - 2. [Performing a PIT join](#2-performing-a-pit-join) - 1. [Early stop sort merge](#21-early-stop-sort-merge) - 2. [Union merge](#22-union-merge) - 3. [Exploding PIT join](#23-exploding-pit-join) + - [Early stop sort merge](#early-stop-sort-merge) - [QuickStart (Scala)](#quickstart-scala) - [Early stop sort merge](#early-stop-sort-merge) - [Union ASOF merge](#union-asof-merge) @@ -29,7 +24,7 @@ The thesis that this project laid the foundation for can be found here: http://w | Dependency | Version | | --------------- | ------- | -| Spark & PySpark | 3.2 | +| Spark & PySpark | 3.5.0 | | Scala | 2.12 | | Python | >=3.6 | @@ -45,6 +40,9 @@ To make the artifacts available for the executors, set the configuration propert Alternatively, set the configuration property `spark.jars` to include the path to the jar-file to make it available for both the driver and executors. +Additionally set `spark.sql.extensions` to include `io.github.ackuq.pit.SparkPIT`. This config +is a comma separated string. + ### Python (PySpark) Configure Spark using the instructions as observed in [the previous section](#scala-spark). @@ -57,72 +55,28 @@ pip install spark-pit ## Quickstart (Python) -### 1. Creating the context - -The object `PitContext` is the entrypoint for all of the functionality of the lirary. You can initialize this context with the following code: - -```py -from ackuq.pit import PitContext - -spark = -pit_context = PitContext(spark) -``` - -### 2. Performing a PIT join - -There are currently 3 ways of executing a PIT join, using an early stop sort merge, union merge algorithm, or with exploding intermediate tables. - -#### 2.1. Early stop sort merge - -```py -pit_join = df1.join(df2, pit_context.pit_udf(df1.ts, df2.ts) & (df1.id == df2.id)) -``` - -#### 2.2. Union merge +### Early stop sort merge ```py -pit_join = pit_context.union( - left=df1, - right=df2, - left_prefix="df1_", - right_prefix="df2_", - left_ts_column = "ts", - right_ts_column = "ts", - partition_cols=["id"], -) -``` +from ackuq.pit.joinPIT import joinPIT -#### 2.3. Exploding PIT join - -```py -pit_join = pit_context.exploding( - left=df1, - right=df2, - left_ts_column=df1["ts"], - right_ts_column=df2["ts"], - partition_cols = [df1["id"], df2["id"]], -) +pit_join = joinPIT(spark, df1, df2, df1.ts, df2.ts, (df1.id == df2.id)) ``` ## Quickstart (Scala) -Instead of using a context, which is done in the Python implementation, all of the functionality is divided into objects. - ### Early stop sort merge ```scala -import io.github.ackuq.pit.EarlyStopSortMerge.{pit, init} -import org.apache.spark.sql.functions.lit +import io.github.ackuq.pit.EarlyStopSortMerge.joinPIT -// Pass the spark session, this will register the required stratergies and optimizer rules. -init(spark) -val pitJoin = df1.join(df2, pit(df1("ts"), df2("ts"), lit(0)) && df1("id") === df2("id")) +val pitJoin = joinPIT(df1, df2, df1("ts"), df2("ts"), df1("id") === df2("id"), "inner", 0) ``` #### Adding tolerance -The UDF takes a third argument (required) for tolerance, when this argument is set to a non-null value, the PIT join does not return matches where the timestamp differ by a specific value. E.g. setting the third argument to `lit(3)` would only accept PIT matches that differ by at most 3 time units. +The final argument is the tolerance, when this argument is set to a non-zero value, the PIT join does not return matches where the timestamps differ by more than the specific value. E.g. setting tolerance to `3` would only accept PIT matches that differ by at most 3 time units. #### Left outer join @@ -131,36 +85,5 @@ The default join type for PIT joins are inner joins, but if you'd like to keep a Usage: ```scala -val pitJoin = df1.join( - df2, pit(df1("ts"), df2("ts"), lit(0)) && df1("id") === df2("id"), - "left" -) -``` - -### Union merge - -```scala -import io.github.ackuq.pit.Union - -val pitJoin = Union.join( - df1, - df2, - leftPrefix = Some("df1_"), - rightPrefix = "df2_", - partitionCols = Seq("id") -) -``` - -### Exploding PIT join - -```scala -import io.github.ackuq.pit.Exploding - -val pitJoin = Exploding.join( - df1, - df2, - leftTSColumn = df1("ts"), - rightTSColumn = df2("ts"), - partitionCols = Seq((df1("id"), df2("id"))) -) +val pitJoin = joinPIT(df1, df2, df1("ts"), df2("ts"), df1("id") === df2("id"), "left") ``` diff --git a/python/README.md b/python/README.md index e18f434..d056734 100644 --- a/python/README.md +++ b/python/README.md @@ -8,76 +8,6 @@ This projects aims to expose different ways of executing PIT-joins, also called In order to run this project in PySpark, you will need to have the JAR file of the Scala implementation be available inside you Spark Session. -## Quickstart - -### 1. Creating the context - -The object `PitContext` is the entrypoint for all of the functionality of the lirary. You can initialize this context with the following code: - -```py -from pyspark import SQLContext -from ackuq.pit import PitContext - -sql_context = SQLContext(spark.sparkContext) -pit_context = PitContext(sql_context) -``` - -### 2. Performing a PIT join - -There are currently 3 ways of executing a PIT join, using an early stop sort merge, union asof algorithm, or with exploding intermediate tables. - -#### 2.1. Early stop sort merge - -```py -pit_join = df1.join(df2, pit_context.pit_udf(df1.ts, df2.ts) & (df1.id == df2.id)) -``` - -##### 2.1.2. Adding tolerance - -In this implementation, tolerance can be added to not allow matches whose timestamp differ by at most some value. To utilize this, set the third argument of the UDF to the desired integer value of the maximum different between two timestamps. - -```py -pit_join = df1.join(df2, pit_context.pit_udf(df1.ts, df2.ts, 100) & (df1.id == df2.id)) -``` - -##### 2.1.3. Left outer join - -Left outer joins are supported in this implementation, the main difference between a regular inner join and a left outer join is that whether or not a left row gets matched with a right row, it will still be a part of the resulting table. In the resulting table, all the left rows that did not find a match have the values of the right columns set to `null`. - -```py -pit_join = df1.join( - df2, - pit_context.pit_udf(df1.ts, df2.ts, 100) & (df1.id == df2.id), - "left" -) -``` - -#### 2.2. Union merge - -```py -pit_join = pit_context.union( - left=df1, - right=df2, - left_prefix="df1_", - right_prefix="df2_", - left_ts_column = "ts", - right_ts_column = "ts", - partition_cols=["id"], -) -``` - -#### 2.3. Exploding PIT join - -```py -pit_join = pit_context.exploding( - left=df1, - right=df2, - left_ts_column=df1["ts"], - right_ts_column=df2["ts"], - partition_cols = [df1["id"], df2["id"]], -) -``` - ## Development ### Testing diff --git a/python/ackuq/pit/joinPIT.py b/python/ackuq/pit/joinPIT.py index 43fd830..442879b 100644 --- a/python/ackuq/pit/joinPIT.py +++ b/python/ackuq/pit/joinPIT.py @@ -22,10 +22,11 @@ # SOFTWARE. # -from pyspark.sql import Column, DataFrame +from pyspark.sql import Column, DataFrame, SparkSession def joinPIT( + spark: SparkSession, left: DataFrame, right: DataFrame, leftPitKey: Column, @@ -34,7 +35,7 @@ def joinPIT( how: str = "inner", tolerance: int = 0, ) -> DataFrame: - jdf = left.sparkSession.sparkContext._jvm.io.github.ackuq.pit.EarlyStopSortMerge.joinPIT( + jdf = spark.sparkContext._jvm.io.github.ackuq.pit.EarlyStopSortMerge.joinPIT( left._jdf, right._jdf, leftPitKey._jc, @@ -43,4 +44,4 @@ def joinPIT( how, tolerance, ) - return DataFrame(jdf, left.sparkSession) + return DataFrame(jdf, spark) diff --git a/python/tests/test_sort_merge.py b/python/tests/test_sort_merge.py index 7c4fe59..a8dce5d 100644 --- a/python/tests/test_sort_merge.py +++ b/python/tests/test_sort_merge.py @@ -37,7 +37,9 @@ def test_two_aligned(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg2 - pit_join = joinPIT(fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"])) + pit_join = joinPIT( + self.spark, fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]) + ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_2.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_2.collect()) @@ -46,7 +48,9 @@ def test_two_misaligned(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = joinPIT(fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"])) + pit_join = joinPIT( + self.spark, fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]) + ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3.collect()) @@ -57,6 +61,7 @@ def test_three_misaligned(self): fg3 = self.small_data.fg3 left = joinPIT( + self.spark, fg1, fg2, fg1["ts"], @@ -64,7 +69,9 @@ def test_three_misaligned(self): (fg1["id"] == fg2["id"]), ) - pit_join = joinPIT(left, fg3, fg1["ts"], fg3["ts"], (fg1["id"] == fg3["id"])) + pit_join = joinPIT( + self.spark, left, fg3, fg1["ts"], fg3["ts"], (fg1["id"] == fg3["id"]) + ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_2_3.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_2_3.collect()) @@ -74,7 +81,13 @@ def test_two_tolerance(self): fg2 = self.small_data.fg3 pit_join = joinPIT( - fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]), tolerance=1 + self.spark, + fg1, + fg2, + fg1["ts"], + fg2["ts"], + (fg1["id"] == fg2["id"]), + tolerance=1, ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3_T1.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3_T1.collect()) @@ -84,7 +97,14 @@ def test_two_tolerance_outer(self): fg2 = self.small_data.fg3 pit_join = joinPIT( - fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]), "left", 1 + self.spark, + fg1, + fg2, + fg1["ts"], + fg2["ts"], + (fg1["id"] == fg2["id"]), + "left", + 1, ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3_T1_OUTER.schema) self.assertEqual(pit_join.collect(), self.small_data.PIT_1_3_T1_OUTER.collect()) From 452dee7fda1d90fe43d9aa06d05ba08e841f67f6 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Thu, 18 Jan 2024 13:19:04 +0000 Subject: [PATCH 11/19] Bump versions (#13) --- python/VERSION | 2 +- scala/VERSION | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/VERSION b/python/VERSION index ab8684e..cc6ce19 100644 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.6.0-tom +0.6.1-tom diff --git a/scala/VERSION b/scala/VERSION index ab8684e..cc6ce19 100644 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.6.0-tom +0.6.1-tom From 8fa2997bc1518b5c5b18bc85ffc4c40fe49a52cd Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Wed, 14 Aug 2024 13:39:46 +0100 Subject: [PATCH 12/19] Avoid checking in proprietary jars --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 037ae54..b90d081 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,8 @@ scala/project/target/ scala/project/project/target/ scala/target/stream/* scala/src/test/scala/Playground.scala +# Avoid checking in proprietary jars +scala/lib .bsp # ignore build files From 55738ce3478de1147cf57b6488fd4e1ade3400d5 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Thu, 14 Nov 2024 18:02:59 +0000 Subject: [PATCH 13/19] Upgrade to spark 3.5.3 (#18) * Use updated spark * Upgrade other deps * Update python side spark dependency * Bump the version * Add wheel to requirements * version bump python side too --- python/VERSION | 2 +- python/requirements_common.txt | 2 +- python/requirements_dev.txt | 3 ++- scala/VERSION | 2 +- scala/build.sbt | 10 +++++----- 5 files changed, 10 insertions(+), 9 deletions(-) diff --git a/python/VERSION b/python/VERSION index cc6ce19..af793f3 100644 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.6.1-tom +0.6.2-tom diff --git a/python/requirements_common.txt b/python/requirements_common.txt index 67e0899..fda0822 100644 --- a/python/requirements_common.txt +++ b/python/requirements_common.txt @@ -1,3 +1,3 @@ # Common requirements used for all environments -pyspark==3.5.0 +pyspark==3.5.3 diff --git a/python/requirements_dev.txt b/python/requirements_dev.txt index 3082dec..a4ef9aa 100644 --- a/python/requirements_dev.txt +++ b/python/requirements_dev.txt @@ -3,4 +3,5 @@ black==22.12.0 flake8==6.0.0 -isort==5.12.0 \ No newline at end of file +isort==5.12.0 +wheel diff --git a/scala/VERSION b/scala/VERSION index cc6ce19..af793f3 100644 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.6.1-tom +0.6.2-tom diff --git a/scala/build.sbt b/scala/build.sbt index d034eec..6335dc8 100644 --- a/scala/build.sbt +++ b/scala/build.sbt @@ -3,7 +3,7 @@ val versionNumber = IO.readLines(new File(versionNumberFile)) ThisBuild / version := versionNumber.head -ThisBuild / scalaVersion := "2.12.15" +ThisBuild / scalaVersion := "2.12.20" lazy val root = (project in file(".")) .settings( @@ -12,8 +12,8 @@ lazy val root = (project in file(".")) ) libraryDependencies ++= Seq( - "org.apache.spark" %% "spark-core" % "3.5.0" % "provided", - "org.apache.spark" %% "spark-sql" % "3.5.0" % "provided", - "org.scalactic" %% "scalactic" % "3.2.12", - "org.scalatest" %% "scalatest" % "3.2.12" % "test" + "org.apache.spark" %% "spark-core" % "3.5.3" % "provided", + "org.apache.spark" %% "spark-sql" % "3.5.3" % "provided", + "org.scalactic" %% "scalactic" % "3.2.19", + "org.scalatest" %% "scalatest" % "3.2.19" % "test" ) From 3ce72981ace41cb58d8674a1ea5dec3169fb0d11 Mon Sep 17 00:00:00 2001 From: Marina Gonzalez Date: Mon, 10 Feb 2025 13:59:14 +0000 Subject: [PATCH 14/19] chore: bump versions for spark 4.0.0rc1 (#19) --- python/VERSION | 2 +- python/requirements_common.txt | 2 +- scala/VERSION | 2 +- scala/build.sbt | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/VERSION b/python/VERSION index af793f3..81fb105 100644 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.6.2-tom +0.7.0-tom diff --git a/python/requirements_common.txt b/python/requirements_common.txt index fda0822..7b149eb 100644 --- a/python/requirements_common.txt +++ b/python/requirements_common.txt @@ -1,3 +1,3 @@ # Common requirements used for all environments -pyspark==3.5.3 +pyspark==4.0.0rc1 diff --git a/scala/VERSION b/scala/VERSION index af793f3..81fb105 100644 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.6.2-tom +0.7.0-tom diff --git a/scala/build.sbt b/scala/build.sbt index 6335dc8..318db5b 100644 --- a/scala/build.sbt +++ b/scala/build.sbt @@ -3,7 +3,7 @@ val versionNumber = IO.readLines(new File(versionNumberFile)) ThisBuild / version := versionNumber.head -ThisBuild / scalaVersion := "2.12.20" +ThisBuild / scalaVersion := "2.13.12" lazy val root = (project in file(".")) .settings( @@ -12,8 +12,8 @@ lazy val root = (project in file(".")) ) libraryDependencies ++= Seq( - "org.apache.spark" %% "spark-core" % "3.5.3" % "provided", - "org.apache.spark" %% "spark-sql" % "3.5.3" % "provided", + "org.apache.spark" %% "spark-core" % "4.0.0-preview1" % "provided", + "org.apache.spark" %% "spark-sql" % "4.0.0-preview1" % "provided", "org.scalactic" %% "scalactic" % "3.2.19", "org.scalatest" %% "scalatest" % "3.2.19" % "test" ) From 957a41138cb8817cb3e43633ea00db468d39bfa8 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 17 Apr 2026 16:00:06 +0100 Subject: [PATCH 15/19] Make comparisons to sort merge join (#21) * Re-order methods to make it more easily comparable to native spark SortMergeJoin * Another comparison comment --- scala/src/main/scala/EarlyStopSortMerge.scala | 2 + .../main/scala/execution/PITJoinExec.scala | 110 ++++++++---------- scala/src/main/scala/logical/PITJoin.scala | 26 +++-- 3 files changed, 67 insertions(+), 71 deletions(-) diff --git a/scala/src/main/scala/EarlyStopSortMerge.scala b/scala/src/main/scala/EarlyStopSortMerge.scala index c20dd03..f1d330e 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, JoinType} import execution.CustomStrategy import logical.PITJoin + object EarlyStopSortMerge { def joinPIT( left: DataFrame, @@ -100,6 +101,7 @@ object EarlyStopSortMerge { tolerance, joinExprs.map(_.expr) ) + // Copying `Dataset.ofRows()`, but using a public constructor for DataFrame (Dataset[Row]). new DataFrame( left.sparkSession, logicalPlan, diff --git a/scala/src/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index 872e6cb..936cb98 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -36,8 +36,8 @@ import org.apache.spark.sql.execution.metric.SQLMetrics import logical.{CustomJoinType, PITJoinType} -/** Performs a PIT join of two child relations. - */ + +// Based on org.apache.spark.sql.execution.joins.SortMergeJoinExec protected[pit] case class PITJoinExec( leftPitKeys: Seq[Expression], rightPitKeys: Seq[Expression], @@ -48,15 +48,7 @@ protected[pit] case class PITJoinExec( right: SparkPlan, returnNulls: Boolean, tolerance: Long -) extends ShuffledJoin - with CodegenSupport { - - override def isSkewJoin: Boolean = false - - override protected def withNewChildrenInternal( - newLeft: SparkPlan, - newRight: SparkPlan - ): PITJoinExec = copy(left = newLeft, right = newRight) +) extends ShuffledJoin { override lazy val metrics: Map[String, SQLMetric] = Map( "numOutputRows" -> SQLMetrics.createMetric( @@ -65,38 +57,6 @@ protected[pit] case class PITJoinExec( ) ) - override def nodeName: String = { - super.nodeName - } - - override def stringArgs: Iterator[Any] = - super.stringArgs.toSeq.dropRight(1).iterator - - override def requiredChildDistribution: Seq[Distribution] = { - if (leftEquiKeys.isEmpty || rightEquiKeys.isEmpty) { - // TODO: This should be improved, but for now just keep everything in one partition - AllTuples :: AllTuples :: Nil - } else { - ClusteredDistribution(leftEquiKeys) :: ClusteredDistribution( - rightEquiKeys - ) :: Nil - } - } - - val customJoinType: CustomJoinType = PITJoinType - - // Set as inner - override def joinType: JoinType = Inner - - override def outputPartitioning: Partitioning = customJoinType match { - // Left and right output partitioning should equal in the results - case PITJoinType => left.outputPartitioning - case x => - throw new IllegalArgumentException( - s"${getClass.getSimpleName} not take $x as the JoinType" - ) - } - override def outputOrdering: Seq[SortOrder] = customJoinType match { // For PIT join, the order should be in descending time order for both sides case PITJoinType => getKeyOrdering(leftKeys, left.outputOrdering) @@ -106,9 +66,7 @@ protected[pit] case class PITJoinExec( ) } - override def leftKeys: Seq[Expression] = leftEquiKeys ++ leftPitKeys - - /** The utility method to get output ordering for left or right side of the + /** The utility method to get output ordering for left or right side of the * join. * * Returns the required ordering for left or right child if @@ -125,7 +83,7 @@ protected[pit] case class PITJoinExec( if (SortOrder.orderingSatisfies(childOutputOrdering, requiredOrdering)) { keys.zip(childOutputOrdering).map { case (key, childOrder) => val sameOrderExpressionsSet = ExpressionSet(childOrder.children) - key - // Changed to descending + // Compared to SortMergeJoinExec the only difference is using descending order. SortOrder(key, Descending, sameOrderExpressionsSet.toSeq) } } else { @@ -133,14 +91,48 @@ protected[pit] case class PITJoinExec( } } + override def requiredChildDistribution: Seq[Distribution] = { + if (leftEquiKeys.isEmpty || rightEquiKeys.isEmpty) { + // TODO: This should be improved, but for now just keep everything in one partition + AllTuples :: AllTuples :: Nil + } else { + ClusteredDistribution(leftEquiKeys) :: ClusteredDistribution( + rightEquiKeys + ) :: Nil + } + } + + override def requiredChildOrdering: Seq[Seq[SortOrder]] = + requiredOrders(leftKeys) :: requiredOrders(rightKeys) :: Nil + private def requiredOrders(keys: Seq[Expression]): Seq[SortOrder] = { // This must be descending in order to agree with the `keyOrdering` defined in `doExecute()`. keys.map(SortOrder(_, Descending)) } - override def requiredChildOrdering: Seq[Seq[SortOrder]] = - requiredOrders(leftKeys) :: requiredOrders(rightKeys) :: Nil + override def isSkewJoin: Boolean = false + + override def nodeName: String = { + super.nodeName + } + + override def stringArgs: Iterator[Any] = + super.stringArgs.toSeq.dropRight(1).iterator + + + val customJoinType: CustomJoinType = PITJoinType + override def joinType: JoinType = Inner // Need to set something valid, for class compatibility + + override def outputPartitioning: Partitioning = customJoinType match { + // Left and right output partitioning should equal in the results + case PITJoinType => left.outputPartitioning + case x => + throw new IllegalArgumentException( + s"${getClass.getSimpleName} not take $x as the JoinType" + ) + } + override def leftKeys: Seq[Expression] = leftEquiKeys ++ leftPitKeys override def rightKeys: Seq[Expression] = rightEquiKeys ++ rightPitKeys protected override def doExecute(): RDD[InternalRow] = { @@ -238,21 +230,14 @@ protected[pit] case class PITJoinExec( } } - /** These generator are for generating for the equi-keys - */ - private def createLeftEquiKeyGenerator(): Projection = UnsafeProjection.create(leftEquiKeys, left.output) private def createRightEquiKeyGenerator(): Projection = UnsafeProjection.create(rightEquiKeys, right.output) - /** These generator are for generating the PIT keys - */ - - private def createLeftPITKeyGenerator(): Projection = { + private def createLeftPITKeyGenerator(): Projection = UnsafeProjection.create(leftPitKeys, left.output) - } private def createRightPITKeyGenerator(): Projection = UnsafeProjection.create(rightPitKeys, right.output) @@ -531,7 +516,7 @@ protected[pit] case class PITJoinExec( | } else if ($tolerance > 0 && ${genToleranceConditions( leftPITKeyVars, rightPITKeyVars - )}) { + )}) { | if($returnNulls) { | $matched = null; | return false; @@ -553,8 +538,6 @@ protected[pit] case class PITJoinExec( (leftRow, matched) } - override def needCopyResult: Boolean = true - override protected def doProduce(ctx: CodegenContext): String = { // Inline mutable state since not many join operations in a task val leftInput = ctx.addMutableState( @@ -609,6 +592,13 @@ protected[pit] case class PITJoinExec( |""".stripMargin } } + + override def needCopyResult: Boolean = true + + override protected def withNewChildrenInternal( + newLeft: SparkPlan, + newRight: SparkPlan + ): PITJoinExec = copy(left = newLeft, right = newRight) } /** Helper class that is used to implement [[PITJoinExec]]. diff --git a/scala/src/main/scala/logical/PITJoin.scala b/scala/src/main/scala/logical/PITJoin.scala index 93189e2..833d83d 100644 --- a/scala/src/main/scala/logical/PITJoin.scala +++ b/scala/src/main/scala/logical/PITJoin.scala @@ -42,6 +42,7 @@ protected[pit] case object PITJoinType extends CustomJoinType { override def sql: String = "PIT" } +// Based on org.apache.spark.sql.catalyst.plans.logical.Join protected[pit] case class PITJoin( left: LogicalPlan, right: LogicalPlan, @@ -53,17 +54,6 @@ protected[pit] case class PITJoin( ) extends BinaryNode with PredicateHelper { - // Joins are only resolved if they don't introduce ambiguous expression ids. - override lazy val resolved: Boolean = { - childrenResolved && - expressions.forall(_.resolved) && - duplicateResolved && - condition.forall(_.dataType == BooleanType) - } - override protected lazy val validConstraints: ExpressionSet = { - left.constraints - } - override def maxRows: Option[Long] = { left.maxRows } @@ -80,9 +70,23 @@ protected[pit] case class PITJoin( children.flatMap(_.metadataOutput) } + + override protected lazy val validConstraints: ExpressionSet = { + left.constraints + } + def duplicateResolved: Boolean = left.outputSet.intersect(right.outputSet).isEmpty + + // Joins are only resolved if they don't introduce ambiguous expression ids. + override lazy val resolved: Boolean = { + childrenResolved && + expressions.forall(_.resolved) && + duplicateResolved && + condition.forall(_.dataType == BooleanType) + } + override protected def withNewChildrenInternal( newLeft: LogicalPlan, newRight: LogicalPlan From d7e7e04eb907ec0a20ddcd80798df0f81427da8f Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 17 Apr 2026 16:02:53 +0100 Subject: [PATCH 16/19] Revert "chore: bump versions for spark 4.0.0rc1 (#19)" (#22) This reverts commit 3ce72981ace41cb58d8674a1ea5dec3169fb0d11. --- python/VERSION | 2 +- python/requirements_common.txt | 2 +- scala/VERSION | 2 +- scala/build.sbt | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/VERSION b/python/VERSION index 81fb105..af793f3 100644 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.7.0-tom +0.6.2-tom diff --git a/python/requirements_common.txt b/python/requirements_common.txt index 7b149eb..fda0822 100644 --- a/python/requirements_common.txt +++ b/python/requirements_common.txt @@ -1,3 +1,3 @@ # Common requirements used for all environments -pyspark==4.0.0rc1 +pyspark==3.5.3 diff --git a/scala/VERSION b/scala/VERSION index 81fb105..af793f3 100644 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.7.0-tom +0.6.2-tom diff --git a/scala/build.sbt b/scala/build.sbt index 318db5b..6335dc8 100644 --- a/scala/build.sbt +++ b/scala/build.sbt @@ -3,7 +3,7 @@ val versionNumber = IO.readLines(new File(versionNumberFile)) ThisBuild / version := versionNumber.head -ThisBuild / scalaVersion := "2.13.12" +ThisBuild / scalaVersion := "2.12.20" lazy val root = (project in file(".")) .settings( @@ -12,8 +12,8 @@ lazy val root = (project in file(".")) ) libraryDependencies ++= Seq( - "org.apache.spark" %% "spark-core" % "4.0.0-preview1" % "provided", - "org.apache.spark" %% "spark-sql" % "4.0.0-preview1" % "provided", + "org.apache.spark" %% "spark-core" % "3.5.3" % "provided", + "org.apache.spark" %% "spark-sql" % "3.5.3" % "provided", "org.scalactic" %% "scalactic" % "3.2.19", "org.scalatest" %% "scalatest" % "3.2.19" % "test" ) From c0b83109132c7ed709b2bdf5d8afc95415456d99 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 17 Apr 2026 16:17:28 +0100 Subject: [PATCH 17/19] Upgrade to spark 4.1.1 (#23) * Bump versions * Working build * Use the row encoder as per up to date spark * All tests pass * Bump python spark version * Update README and make it 0.8.0 --- README.md | 4 ++-- python/VERSION | 2 +- python/requirements_common.txt | 2 +- scala/.scalafmt.conf | 2 +- scala/VERSION | 2 +- scala/build.sbt | 6 +++--- scala/src/main/scala/EarlyStopSortMerge.scala | 17 ++++++++++------- .../main/scala/execution/CustomStrategy.scala | 2 +- scala/src/test/scala/EarlyStopMergeTests.scala | 2 +- .../test/scala/SparkSessionTestWrapper.scala | 2 +- scala/src/test/scala/data/SmallData.scala | 4 ++-- .../test/scala/data/SmallDataExploding.scala | 3 ++- .../test/scala/data/SmallDataSortMerge.scala | 4 ++-- scala/src/test/scala/data/SmallDataUnion.scala | 3 ++- 14 files changed, 30 insertions(+), 25 deletions(-) mode change 100644 => 100755 python/VERSION mode change 100644 => 100755 scala/VERSION mode change 100644 => 100755 scala/build.sbt diff --git a/README.md b/README.md index aafa5ac..5ff8a00 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,8 @@ The thesis that this project laid the foundation for can be found here: http://w | Dependency | Version | | --------------- | ------- | -| Spark & PySpark | 3.5.0 | -| Scala | 2.12 | +| Spark & PySpark | 4.1.1 | +| Scala | 2.13 | | Python | >=3.6 | ## Installation diff --git a/python/VERSION b/python/VERSION old mode 100644 new mode 100755 index af793f3..844f6a9 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.6.2-tom +0.6.3 diff --git a/python/requirements_common.txt b/python/requirements_common.txt index fda0822..6f4e83d 100644 --- a/python/requirements_common.txt +++ b/python/requirements_common.txt @@ -1,3 +1,3 @@ # Common requirements used for all environments -pyspark==3.5.3 +pyspark==4.1.1 diff --git a/scala/.scalafmt.conf b/scala/.scalafmt.conf index 7fee364..49c2132 100644 --- a/scala/.scalafmt.conf +++ b/scala/.scalafmt.conf @@ -1,3 +1,3 @@ version = 3.4.3 -runner.dialect = scala212 \ No newline at end of file +runner.dialect = scala213 \ No newline at end of file diff --git a/scala/VERSION b/scala/VERSION old mode 100644 new mode 100755 index af793f3..a3df0a6 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.6.2-tom +0.8.0 diff --git a/scala/build.sbt b/scala/build.sbt old mode 100644 new mode 100755 index 6335dc8..da4f694 --- a/scala/build.sbt +++ b/scala/build.sbt @@ -3,7 +3,7 @@ val versionNumber = IO.readLines(new File(versionNumberFile)) ThisBuild / version := versionNumber.head -ThisBuild / scalaVersion := "2.12.20" +ThisBuild / scalaVersion := "2.13.18" lazy val root = (project in file(".")) .settings( @@ -12,8 +12,8 @@ lazy val root = (project in file(".")) ) libraryDependencies ++= Seq( - "org.apache.spark" %% "spark-core" % "3.5.3" % "provided", - "org.apache.spark" %% "spark-sql" % "3.5.3" % "provided", + "org.apache.spark" %% "spark-core" % "4.1.1" % "provided", + "org.apache.spark" %% "spark-sql" % "4.1.1" % "provided", "org.scalactic" %% "scalactic" % "3.2.19", "org.scalatest" %% "scalatest" % "3.2.19" % "test" ) diff --git a/scala/src/main/scala/EarlyStopSortMerge.scala b/scala/src/main/scala/EarlyStopSortMerge.scala index f1d330e..450c1f8 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -26,11 +26,11 @@ package io.github.ackuq.pit import org.apache.spark.sql.{ Column, - DataFrame, SparkSessionExtensionsProvider, SparkSessionExtensions } -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.catalyst.encoders.RowEncoder import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, JoinType} import execution.CustomStrategy @@ -92,20 +92,23 @@ object EarlyStopSortMerge { ) } + val sparkSession = left.sparkSession + def toExpression(column: Column) = sparkSession.expression(column) + val logicalPlan = PITJoin( left.queryExecution.analyzed, right.queryExecution.analyzed, - leftPitExpression.expr, - rightPitExpression.expr, + toExpression(leftPitExpression), + toExpression(rightPitExpression), parsedJoinType == LeftOuter, tolerance, - joinExprs.map(_.expr) + joinExprs.map(toExpression(_)) ) // Copying `Dataset.ofRows()`, but using a public constructor for DataFrame (Dataset[Row]). new DataFrame( - left.sparkSession, + sparkSession, logicalPlan, - ExpressionEncoder(logicalPlan.schema) + RowEncoder.encoderFor(logicalPlan.schema) ) } } diff --git a/scala/src/main/scala/execution/CustomStrategy.scala b/scala/src/main/scala/execution/CustomStrategy.scala index f783ce1..08ce0ad 100644 --- a/scala/src/main/scala/execution/CustomStrategy.scala +++ b/scala/src/main/scala/execution/CustomStrategy.scala @@ -25,7 +25,7 @@ package io.github.ackuq.pit package execution -import org.apache.spark.sql.Strategy +import org.apache.spark.sql.classic.Strategy import org.apache.spark.sql.catalyst.expressions.{PredicateHelper, RowOrdering} import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index ecbc0ca..cef6eac 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -24,7 +24,7 @@ package io.github.ackuq.pit -import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.StructType diff --git a/scala/src/test/scala/SparkSessionTestWrapper.scala b/scala/src/test/scala/SparkSessionTestWrapper.scala index c6b6d6b..5259cfa 100644 --- a/scala/src/test/scala/SparkSessionTestWrapper.scala +++ b/scala/src/test/scala/SparkSessionTestWrapper.scala @@ -24,7 +24,7 @@ package io.github.ackuq.pit -import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.classic.SparkSession trait SparkSessionTestWrapper { val spark: SparkSession = SparkSession diff --git a/scala/src/test/scala/data/SmallData.scala b/scala/src/test/scala/data/SmallData.scala index 3222012..c463591 100644 --- a/scala/src/test/scala/data/SmallData.scala +++ b/scala/src/test/scala/data/SmallData.scala @@ -25,9 +25,9 @@ package io.github.ackuq.pit package data -import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.Row -import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.types.StringType import org.apache.spark.sql.types.StructField diff --git a/scala/src/test/scala/data/SmallDataExploding.scala b/scala/src/test/scala/data/SmallDataExploding.scala index cc395df..a66497d 100644 --- a/scala/src/test/scala/data/SmallDataExploding.scala +++ b/scala/src/test/scala/data/SmallDataExploding.scala @@ -31,7 +31,8 @@ import org.apache.spark.sql.types.{ StructField, StructType } -import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.Row class SmallDataExploding(spark: SparkSession) extends SmallData(spark) { private val PIT_1_2_RAW = Seq( diff --git a/scala/src/test/scala/data/SmallDataSortMerge.scala b/scala/src/test/scala/data/SmallDataSortMerge.scala index 2eeed23..2e7622b 100644 --- a/scala/src/test/scala/data/SmallDataSortMerge.scala +++ b/scala/src/test/scala/data/SmallDataSortMerge.scala @@ -25,9 +25,9 @@ package io.github.ackuq.pit package data -import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.Row -import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.types.StringType import org.apache.spark.sql.types.StructField diff --git a/scala/src/test/scala/data/SmallDataUnion.scala b/scala/src/test/scala/data/SmallDataUnion.scala index 7708e58..111489e 100644 --- a/scala/src/test/scala/data/SmallDataUnion.scala +++ b/scala/src/test/scala/data/SmallDataUnion.scala @@ -31,7 +31,8 @@ import org.apache.spark.sql.types.{ StructField, StructType } -import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.Row class SmallDataUnion(spark: SparkSession) extends SmallData(spark) { private val PIT_1_2_RAW = Seq( From 2e13230a95fcc2937e4205888d51ecf6a657c26e Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 17 Apr 2026 16:23:41 +0100 Subject: [PATCH 18/19] Bump python version (#24) --- python/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/VERSION b/python/VERSION index 844f6a9..a3df0a6 100755 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.6.3 +0.8.0 From ff5dc9bcff79a8768ae13e2f7465e33cc9528bac Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Tue, 5 May 2026 15:32:29 +0100 Subject: [PATCH 19/19] Fail fast on invalid time column (#25) * Initial test and implementation * Use column expression instead of modifying dataframe column * Update implementation based on an example from inside spark * Bump versions * Remove unused import * Unabbreviate * Switch to AnalysisException --- python/VERSION | 2 +- scala/VERSION | 2 +- scala/src/main/scala/EarlyStopSortMerge.scala | 60 ++++++++++++------- .../src/test/scala/EarlyStopMergeTests.scala | 37 +++++++++++- 4 files changed, 77 insertions(+), 24 deletions(-) diff --git a/python/VERSION b/python/VERSION index a3df0a6..6f4eebd 100755 --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.8.0 +0.8.1 diff --git a/scala/VERSION b/scala/VERSION index a3df0a6..6f4eebd 100755 --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.8.0 +0.8.1 diff --git a/scala/src/main/scala/EarlyStopSortMerge.scala b/scala/src/main/scala/EarlyStopSortMerge.scala index 450c1f8..b989b56 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -24,15 +24,16 @@ package io.github.ackuq.pit -import org.apache.spark.sql.{ - Column, - SparkSessionExtensionsProvider, - SparkSessionExtensions -} -import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.Column +import org.apache.spark.sql.SparkSessionExtensions +import org.apache.spark.sql.SparkSessionExtensionsProvider import org.apache.spark.sql.catalyst.encoders.RowEncoder -import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, JoinType} - +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.JoinType +import org.apache.spark.sql.catalyst.plans.LeftOuter +import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.types.NumericType import execution.CustomStrategy import logical.PITJoin @@ -41,15 +42,15 @@ object EarlyStopSortMerge { def joinPIT( left: DataFrame, right: DataFrame, - leftPitExpression: Column, - rightPitExpression: Column, + leftPitColumn: Column, + rightPitColumn: Column, joinType: String, tolerance: Long ): DataFrame = joinPIT( left, right, - leftPitExpression, - rightPitExpression, + leftPitColumn, + rightPitColumn, None, joinType, tolerance @@ -58,16 +59,16 @@ object EarlyStopSortMerge { def joinPIT( left: DataFrame, right: DataFrame, - leftPitExpression: Column, - rightPitExpression: Column, + leftPitColumn: Column, + rightPitColumn: Column, joinExprs: Column, joinType: String, tolerance: Long ): DataFrame = joinPIT( left, right, - leftPitExpression, - rightPitExpression, + leftPitColumn, + rightPitColumn, Some(joinExprs), joinType, tolerance @@ -76,8 +77,8 @@ object EarlyStopSortMerge { def joinPIT( left: DataFrame, right: DataFrame, - leftPitExpression: Column, - rightPitExpression: Column, + leftPitColumn: Column, + rightPitColumn: Column, joinExprs: Option[Column], joinType: String, tolerance: Long @@ -92,14 +93,33 @@ object EarlyStopSortMerge { ) } + val sparkSession = left.sparkSession def toExpression(column: Column) = sparkSession.expression(column) + val leftPitExpression = toExpression(leftPitColumn) + val rightPitExpression = toExpression(rightPitColumn) + + Seq("left" -> leftPitExpression.dataType, "right" -> rightPitExpression.dataType).foreach { + case (side, dataType) => + if (!dataType.isInstanceOf[NumericType]) { + throw new AnalysisException( + message = s"PIT key on $side side must be a numeric type, got $dataType", + line = None, + startPosition = None, + cause = None, + errorClass = None, + messageParameters = Map.empty, + context = Array.empty + ) + } + } + val logicalPlan = PITJoin( left.queryExecution.analyzed, right.queryExecution.analyzed, - toExpression(leftPitExpression), - toExpression(rightPitExpression), + leftPitExpression, + rightPitExpression, parsedJoinType == LeftOuter, tolerance, joinExprs.map(toExpression(_)) diff --git a/scala/src/test/scala/EarlyStopMergeTests.scala b/scala/src/test/scala/EarlyStopMergeTests.scala index cef6eac..f5131e4 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -24,15 +24,16 @@ package io.github.ackuq.pit -import org.apache.spark.sql.classic.DataFrame +import io.github.ackuq.pit.execution.CustomStrategy +import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode +import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.StructType import org.scalatest.flatspec.AnyFlatSpec import EarlyStopSortMerge.joinPIT import data.SmallDataSortMerge -import io.github.ackuq.pit.execution.CustomStrategy class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { val smallData = new SmallDataSortMerge(spark) @@ -461,6 +462,38 @@ class EarlyStopMergeTests extends AnyFlatSpec with SparkSessionTestWrapper { testJoiningThreeDataframes("left", smallData.PIT_3_OUTER_schema) } + def testFailNonNumericPITKeys() { + val fg1 = smallData.fg1 + val fg2 = smallData.fg2 + + intercept[AnalysisException] { + val pitJoin = joinPIT( + fg1, + fg2, + fg1("ts"), + fg2("ts").cast("string"), + fg1("id") === fg2("id"), + "inner", + 0, + ) + } + intercept[AnalysisException] { + val pitJoin = joinPIT( + fg1, + fg2, + fg1("ts").cast("string"), + fg2("ts"), + fg1("id") === fg2("id"), + "inner", + 0, + ) + } + } + + testBothCodegenAndInterpreted("fail_for_non_long_pit_key_type") { + testFailNonNumericPITKeys() + } + testBothCodegenAndInterpreted("fail_during_planning_for_non_equi_condition") { val fg1 = smallData.fg1 val fg2 = smallData.fg2