diff --git a/.gitignore b/.gitignore index 600ab55..b90d081 100644 --- a/.gitignore +++ b/.gitignore @@ -2,13 +2,18 @@ .idea *.iml .vscode +.metals +.bloop # Scala related files +scala/project scala/target 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 @@ -16,4 +21,4 @@ scala/src/test/scala/Playground.scala **/dist # Mac files -.DS_Store \ No newline at end of file +.DS_Store diff --git a/README.md b/README.md index e2ccbe9..5ff8a00 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,8 +24,8 @@ The thesis that this project laid the foundation for can be found here: http://w | Dependency | Version | | --------------- | ------- | -| Spark & PySpark | 3.2 | -| Scala | 2.12 | +| Spark & PySpark | 4.1.1 | +| Scala | 2.13 | | Python | >=3.6 | ## Installation @@ -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,73 +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 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 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 @@ -132,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/VERSION b/python/VERSION old mode 100644 new mode 100755 index 79a2734..6f4eebd --- a/python/VERSION +++ b/python/VERSION @@ -1 +1 @@ -0.5.0 \ No newline at end of file +0.8.1 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 000f3c4..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, SQLContext -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._sql_context.sparkSession._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, sql_context: SQLContext) -> None: - self._sql_context = sql_context - self._sc = self._sql_context._sc # type: ignore - 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..442879b --- /dev/null +++ b/python/ackuq/pit/joinPIT.py @@ -0,0 +1,47 @@ +# +# 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, SparkSession + + +def joinPIT( + spark: SparkSession, + left: DataFrame, + right: DataFrame, + leftPitKey: Column, + rightPitKey: Column, + on: Column, + how: str = "inner", + tolerance: int = 0, +) -> DataFrame: + jdf = spark.sparkContext._jvm.io.github.ackuq.pit.EarlyStopSortMerge.joinPIT( + left._jdf, + right._jdf, + leftPitKey._jc, + rightPitKey._jc, + on._jc, + how, + tolerance, + ) + return DataFrame(jdf, spark) diff --git a/python/requirements_common.txt b/python/requirements_common.txt index 3f8a5ac..6f4e83d 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==4.1.1 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/python/tests/test_sort_merge.py b/python/tests/test_sort_merge.py index df29a09..a8dce5d 100644 --- a/python/tests/test_sort_merge.py +++ b/python/tests/test_sort_merge.py @@ -22,22 +22,23 @@ # 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( + self.spark, fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]) ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_2.schema) @@ -47,9 +48,8 @@ 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( + self.spark, fg1, fg2, fg1["ts"], fg2["ts"], (fg1["id"] == fg2["id"]) ) self.assertSchemaEqual(pit_join.schema, self.small_data.PIT_1_3.schema) @@ -60,14 +60,17 @@ def test_three_misaligned(self): fg2 = self.small_data.fg2 fg3 = self.small_data.fg3 - left = fg1.join( + left = joinPIT( + self.spark, + 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( + 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) @@ -77,10 +80,14 @@ def test_two_tolerance(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = fg1.join( + pit_join = joinPIT( + self.spark, + fg1, fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"], 1) - & (fg1["id"] == fg2["id"]), + 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 +96,15 @@ def test_two_tolerance_outer(self): fg1 = self.small_data.fg1 fg2 = self.small_data.fg3 - pit_join = fg1.join( + pit_join = joinPIT( + self.spark, + fg1, fg2, - self.pit_context.pit_udf(fg1["ts"], fg2["ts"], 1) - & (fg1["id"] == fg2["id"]), + 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 0387822..8d7d1d8 100644 --- a/python/tests/utils.py +++ b/python/tests/utils.py @@ -25,30 +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.sql_context = SQLContext(self.spark.sparkContext) - self.pit_context = PitContext(self.sql_context) - 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/.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 5d4294b..6f4eebd --- a/scala/VERSION +++ b/scala/VERSION @@ -1 +1 @@ -0.5.1 \ No newline at end of file +0.8.1 diff --git a/scala/build.sbt b/scala/build.sbt old mode 100644 new mode 100755 index 74a51ac..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.15" +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.2.1" % "provided", - "org.apache.spark" %% "spark-sql" % "3.2.1" % "provided", - "org.scalactic" %% "scalactic" % "3.2.12", - "org.scalatest" %% "scalatest" % "3.2.12" % "test" + "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 5f439e5..b989b56 100644 --- a/scala/src/main/scala/EarlyStopSortMerge.scala +++ b/scala/src/main/scala/EarlyStopSortMerge.scala @@ -24,36 +24,117 @@ package io.github.ackuq.pit +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 +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.PITRule +import logical.PITJoin -import org.apache.spark.sql.expressions.UserDefinedFunction -import org.apache.spark.sql.functions.udf -import org.apache.spark.sql.{Column, SparkSession} 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 + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitColumn: Column, + rightPitColumn: Column, + joinType: String, + tolerance: Long + ): DataFrame = joinPIT( + left, + right, + leftPitColumn, + rightPitColumn, + None, + joinType, + tolerance + ) + + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitColumn: Column, + rightPitColumn: Column, + joinExprs: Column, + joinType: String, + tolerance: Long + ): DataFrame = joinPIT( + left, + right, + leftPitColumn, + rightPitColumn, + Some(joinExprs), + joinType, + tolerance + ) + + def joinPIT( + left: DataFrame, + right: DataFrame, + leftPitColumn: Column, + rightPitColumn: 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" + ) } - if (!spark.experimental.extraStrategies.contains(PITRule)) { - spark.experimental.extraOptimizations = - spark.experimental.extraOptimizations :+ PITRule + + + 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, + leftPitExpression, + rightPitExpression, + parsedJoinType == LeftOuter, + tolerance, + joinExprs.map(toExpression(_)) + ) + // Copying `Dataset.ofRows()`, but using a public constructor for DataFrame (Dataset[Row]). + new DataFrame( + sparkSession, + logicalPlan, + RowEncoder.encoderFor(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/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/main/scala/execution/PITJoinExec.scala b/scala/src/main/scala/execution/PITJoinExec.scala index c3f14cc..936cb98 100644 --- a/scala/src/main/scala/execution/PITJoinExec.scala +++ b/scala/src/main/scala/execution/PITJoinExec.scala @@ -20,22 +20,24 @@ 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 -/** Performs a PIT join of two child relations. - */ +import logical.{CustomJoinType, PITJoinType} + + +// Based on org.apache.spark.sql.execution.joins.SortMergeJoinExec protected[pit] case class PITJoinExec( leftPitKeys: Seq[Expression], rightPitKeys: Seq[Expression], @@ -46,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( @@ -63,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 { - HashClusteredDistribution(leftEquiKeys) :: HashClusteredDistribution( - 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) @@ -104,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 @@ -123,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 { @@ -131,33 +91,55 @@ 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] = { 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,25 +184,20 @@ protected[pit] case class PITJoinExec( new GenericInternalRow(right.output.length) override def advanceNext(): Boolean = { - while (smjScanner.findNextInnerJoinRows()) { - currentRightMatch = smjScanner.getBufferedMatch - currentLeftRow = smjScanner.getStreamedRow - if (returnNulls && currentRightMatch == null) { + if (smjScanner.findNextJoinRow()) { + currentRightMatch = smjScanner.getRightMatch + currentLeftRow = smjScanner.getLeftRow + if (currentRightMatch == null) { 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 } @@ -253,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) @@ -436,10 +406,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 @@ -471,10 +437,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 |){ @@ -487,12 +475,17 @@ 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; | } | do { | if($rightRow == null) { | if(!rightIter.hasNext()) { + | $matched = null; | return false; | } | $rightRow = (InternalRow) rightIter.next(); @@ -523,7 +516,7 @@ protected[pit] case class PITJoinExec( | } else if ($tolerance > 0 && ${genToleranceConditions( leftPITKeyVars, rightPITKeyVars - )}) { + )}) { | if($returnNulls) { | $matched = null; | return false; @@ -545,34 +538,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 = { // Inline mutable state since not many join operations in a task val leftInput = ctx.addMutableState( @@ -597,49 +562,17 @@ 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();" if (returnNulls) { s""" |while($leftInput.hasNext()) { - | findNextInnerJoinRows($leftInput, $rightInput); + | findNextJoinRows($leftInput, $rightInput); | ${leftVarDecl.mkString("\n")} | ${beforeLoop.trim} | InternalRow $rightRow = (InternalRow) $matched; - | ${condCheck.trim} | $numOutput.add(1); | ${consume(ctx, leftVars ++ rightVars)}; | if (shouldStop()) return; @@ -647,11 +580,10 @@ 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; - | ${condCheck.trim} | $numOutput.add(1); | ${consume(ctx, leftVars ++ rightVars)} | if (shouldStop()) return; @@ -660,218 +592,228 @@ 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]]. * - * 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 findNextInnerJoinRows(): 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 - false - } 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/main/scala/execution/Patterns.scala b/scala/src/main/scala/execution/Patterns.scala index 15cb092..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) @@ -108,34 +108,22 @@ 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)) - 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, - otherPredicates.reduceOption(And), + None, join.returnNulls, join.tolerance, join.left, diff --git a/scala/src/main/scala/logical/PITJoin.scala b/scala/src/main/scala/logical/PITJoin.scala index f4e37d8..833d83d 100644 --- a/scala/src/main/scala/logical/PITJoin.scala +++ b/scala/src/main/scala/logical/PITJoin.scala @@ -42,27 +42,18 @@ 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, - pitCondition: Expression, + leftPitKey: Expression, + rightPitKey: Expression, returnNulls: Boolean, tolerance: Long, condition: Option[Expression] ) 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 } @@ -79,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 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 29a0c05..f5131e4 100644 --- a/scala/src/test/scala/EarlyStopMergeTests.scala +++ b/scala/src/test/scala/EarlyStopMergeTests.scala @@ -24,102 +24,491 @@ package io.github.ackuq.pit -import EarlyStopSortMerge.pit -import data.SmallDataSortMerge - +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 + 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 + } - val pitJoin = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") - ) + 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 = { + + leftDataFrame.show() + rightDataFrame.show() + + val pitJoin = joinPIT( + leftDataFrame, + rightDataFrame, + leftDataFrame("ts"), + rightDataFrame("ts"), + leftDataFrame("id") === rightDataFrame("id"), + joinType, + tolerance, + ) + + pitJoin.explain("codegen") + pitJoin.printSchema() + pitJoin.show() - 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())) + expectedDataFrame.show() + + 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 + ) + } + + 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 + ) + } + + 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_value_nulls, + smallData.fg1_with_value_nulls, + smallData.PIT_3_1_WITH_NULLS, + smallData.PIT_2_NULLABLE_schema, + 0 + ) + } + + testBothCodegenAndInterpreted( + "left_join_preserves_input_nulls" + ) { + testSearchingBackwardForMatches( + "left", + smallData.fg3_with_value_nulls, + smallData.fg1_with_value_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_value_nulls, + smallData.fg1_with_value_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_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 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())) + } - 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_then_filter" + ) { + testJoinThenFilter() } - it should "Perform a PIT join with three dataframes, misaligned timestamps" in { + def testJoiningThreeDataframes( + joinType: String, + expectedSchema: StructType + ): Unit = { val fg1 = smallData.fg1 val fg2 = smallData.fg2 val fg3 = smallData.fg3 - val left = - fg1.join( - fg2, - pit(fg1("ts"), fg2("ts"), lit(0)) && fg1("id") === fg2("id") - ) + 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") - ) + val pitJoin = joinPIT( + left, + fg3, + fg1("ts"), + fg3("ts"), + fg1("id") === fg3("id"), + joinType, + 0, + ) - 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 { + testBothCodegenAndInterpreted("inner_join_three_dataframes") { + testJoiningThreeDataframes("inner", smallData.PIT_3_schema) + } + testBothCodegenAndInterpreted("left_join_three_dataframes") { + testJoiningThreeDataframes("left", smallData.PIT_3_OUTER_schema) + } + + def testFailNonNumericPITKeys() { val fg1 = smallData.fg1 - val fg2 = smallData.fg3 + val fg2 = smallData.fg2 - 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())) + 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, + ) + } } - it should "Be able to perform a left outer PIT join with tolerance, misaligned timestamps" in { + testBothCodegenAndInterpreted("fail_for_non_long_pit_key_type") { + testFailNonNumericPITKeys() + } + + testBothCodegenAndInterpreted("fail_during_planning_for_non_equi_condition") { val fg1 = smallData.fg1 - val fg2 = smallData.fg3 + val fg2 = smallData.fg2 - val pitJoin = fg1.join( + val pitJoin = joinPIT( + fg1, 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())) + 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 9bf7fe0..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 @@ -33,5 +33,8 @@ 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/SmallData.scala b/scala/src/test/scala/data/SmallData.scala index 8dae94d..c463591 100644 --- a/scala/src/test/scala/data/SmallData.scala +++ b/scala/src/test/scala/data/SmallData.scala @@ -25,50 +25,138 @@ 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.classic.DataFrame +import org.apache.spark.sql.Row +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 +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_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"), + Row(1, 6, "f3-1-6"), + Row(2, 8, "f3-2-8"), + Row(1, 10, "f3-1-10") + ) + 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( - 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_value_nullable: StructType = StructType( Seq( StructField("id", IntegerType, nullable = false), StructField("ts", IntegerType, nullable = false), + StructField("value", StringType, nullable = true) + ) + ) + + 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(DATA_RAW.head), schema) + spark.createDataFrame(spark.sparkContext.parallelize(RAW_1), schema) + 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_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(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_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_KEY_NULLS), + schema_keys_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/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 07d80d2..2e7622b 100644 --- a/scala/src/test/scala/data/SmallDataSortMerge.scala +++ b/scala/src/test/scala/data/SmallDataSortMerge.scala @@ -25,49 +25,116 @@ 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.classic.DataFrame +import org.apache.spark.sql.Row +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 +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_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_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"), + Row(1, 6, "f3-1-6", 1, 5, null) + ) + 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) + ) + 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_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_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"), + 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 +144,29 @@ 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_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), StructField("ts", IntegerType, nullable = false), @@ -87,8 +176,29 @@ 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_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), StructField("ts", IntegerType, nullable = false), @@ -101,6 +211,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 @@ -109,6 +232,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), @@ -120,8 +257,53 @@ 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_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 + ) + + val PIT_3_1_T1_WITH_NULLS: DataFrame = spark.createDataFrame( + 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_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_VALUE_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 + ) } 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(