From 591a60df788ae72226375f2d3e85c203200b4b93 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Wed, 11 Sep 2024 13:03:25 -0700 Subject: [PATCH 001/189] [SPARK-49602][BUILD] Fix `assembly/pom.xml` to use `{project.version}` instead of `{version}` ### What changes were proposed in this pull request? This PR aims to fix `assembly/pom.xml` to use `{project.version}` instead of `{version}`. The original change was introduced recently by - #47402 ### Why are the changes needed? **BEFORE** ``` $ mvn clean | head -n9 [INFO] Scanning for projects... [WARNING] [WARNING] Some problems were encountered while building the effective model for org.apache.spark:spark-assembly_2.13:pom:4.0.0-SNAPSHOT [WARNING] The expression ${version} is deprecated. Please use ${project.version} instead. [WARNING] [WARNING] It is highly recommended to fix these problems because they threaten the stability of your build. [WARNING] [WARNING] For this reason, future Maven versions might no longer support building such malformed projects. [WARNING] ``` **AFTER** ``` $ mvn clean | head -n9 [INFO] Scanning for projects... [INFO] ------------------------------------------------------------------------ [INFO] Detecting the operating system and CPU architecture [INFO] ------------------------------------------------------------------------ [INFO] os.detected.name: osx [INFO] os.detected.arch: aarch_64 [INFO] os.detected.version: 15.0 [INFO] os.detected.version.major: 15 [INFO] os.detected.version.minor: 0 ``` ### Does this PR introduce _any_ user-facing change? No, this is a dev-only change for building distribution. ### How was this patch tested? Manual test. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48081 from dongjoon-hyun/SPARK-49602. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- assembly/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/assembly/pom.xml b/assembly/pom.xml index 8b21f7e808ce1..4b074a88dab4a 100644 --- a/assembly/pom.xml +++ b/assembly/pom.xml @@ -200,7 +200,7 @@ cp - ${basedir}/../connector/connect/client/jvm/target/spark-connect-client-jvm_${scala.binary.version}-${version}.jar + ${basedir}/../connector/connect/client/jvm/target/spark-connect-client-jvm_${scala.binary.version}-${project.version}.jar ${basedir}/target/scala-${scala.binary.version}/jars/connect-repl From 07f5b2c1c5ff65df5cf6067bd1108c1cad9dd70d Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 11 Sep 2024 13:18:56 -0700 Subject: [PATCH 002/189] [SPARK-49155][SQL][SS] Use more appropriate parameter type to construct `GenericArrayData` ### What changes were proposed in this pull request? Referring to the test results of `GenericArrayDataBenchmark`, using an Array of Any to construct `GenericArrayData` is more efficient compared to other scenarios: https://github.com/apache/spark/blob/master/sql/catalyst/benchmarks/GenericArrayDataBenchmark-results.txt ``` OpenJDK 64-Bit Server VM 17.0.11+9-LTS on Linux 6.5.0-1018-azure AMD EPYC 7763 64-Core Processor constructor: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ arrayOfAny 6 6 0 1620.1 0.6 1.0X arrayOfAnyAsObject 6 6 0 1620.1 0.6 1.0X arrayOfAnyAsSeq 155 155 1 64.7 15.5 0.0X arrayOfInt 253 254 1 39.6 25.3 0.0X arrayOfIntAsObject 252 253 1 39.7 25.2 0.0X ``` So this pr optimizes some processes of constructing `GenericArrayData` in Spark code: 1. In `ArraysZip#eval` and `XPathList#nullSafeEval`, the originally defined arrays of specific types are changed to data of type `AnyRef` to avoid additional collection copying when constructing `GenericArrayData`. This is because the `Array[AnyRef]` type can also match the `case array: Array[Any] => array` branch in the following code: https://github.com/apache/spark/blob/af70aafd330fdbb6ce0d5b3efbcb180cda488695/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GenericArrayData.scala#L42-L48 2. In `HistogramNumeric#eval`, an `IndexedSeq[InternalRow]` was originally used to construct `GenericArrayData`. Since the length of the collection is known, it can be refactored to use `Array[AnyRef]` to construct `GenericArrayData`. 3. For other cases, when constructing `GenericArrayData`, the current input parameter is `${input}.toArray` now. It is changed to `${input}.toArray[Any]` to avoid another collection copy during the construction of `GenericArrayData`. ### Why are the changes needed? Using an Array of `Any|AnyRef` to construct `GenericArrayData` can improve performance by reducing collection copying. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Pass GitHub Actions ### Was this patch authored or co-authored using generative AI tooling? No Closes #47662 from LuciferYang/GenericArrayData-constructor. Lead-authored-by: yangjie01 Co-authored-by: YangJie Signed-off-by: Dongjoon Hyun --- .../spark/sql/kafka010/KafkaRecordToRowConverter.scala | 2 +- .../expressions/aggregate/HistogramNumeric.scala | 9 +++++---- .../sql/catalyst/expressions/aggregate/collect.scala | 2 +- .../sql/catalyst/expressions/collectionOperations.scala | 2 +- .../spark/sql/catalyst/expressions/jsonExpressions.scala | 2 +- .../spark/sql/catalyst/expressions/xml/xpath.scala | 2 +- .../apache/spark/sql/catalyst/json/JacksonParser.scala | 2 +- .../spark/sql/execution/command/CommandUtils.scala | 2 +- 8 files changed, 12 insertions(+), 11 deletions(-) diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRecordToRowConverter.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRecordToRowConverter.scala index 56456f9b1f776..8d0bcc5816775 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRecordToRowConverter.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaRecordToRowConverter.scala @@ -50,7 +50,7 @@ private[kafka010] class KafkaRecordToRowConverter { new GenericArrayData(cr.headers.iterator().asScala .map(header => InternalRow(UTF8String.fromString(header.key()), header.value()) - ).toArray) + ).toArray[Any]) } else { null } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala index ba26c5a1022d0..eda2c742ab4b5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala @@ -143,7 +143,8 @@ case class HistogramNumeric( if (buffer.getUsedBins < 1) { null } else { - val result = (0 until buffer.getUsedBins).map { index => + val array = new Array[AnyRef](buffer.getUsedBins) + (0 until buffer.getUsedBins).foreach { index => // Note that the 'coord.x' and 'coord.y' have double-precision floating point type here. val coord = buffer.getBin(index) if (propagateInputType) { @@ -163,16 +164,16 @@ case class HistogramNumeric( coord.x.toLong case _ => coord.x } - InternalRow.apply(result, coord.y) + array(index) = InternalRow.apply(result, coord.y) } else { // Otherwise, just apply the double-precision values in 'coord.x' and 'coord.y' to the // output row directly. In this case: 'SELECT histogram_numeric(val, 3) // FROM VALUES (0L), (1L), (2L), (10L) AS tab(col)' returns an array of structs where the // first field has DoubleType. - InternalRow.apply(coord.x, coord.y) + array(index) = InternalRow.apply(coord.x, coord.y) } } - new GenericArrayData(result) + new GenericArrayData(array) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala index e77622a26d90a..c593c8bfb8341 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala @@ -170,7 +170,7 @@ case class CollectSet( override def eval(buffer: mutable.HashSet[Any]): Any = { val array = child.dataType match { case BinaryType => - buffer.iterator.map(_.asInstanceOf[ArrayData].toByteArray()).toArray + buffer.iterator.map(_.asInstanceOf[ArrayData].toByteArray()).toArray[Any] case _ => buffer.toArray } new GenericArrayData(array) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala index 375a2bde59230..5d5aece35383e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala @@ -433,7 +433,7 @@ case class ArraysZip(children: Seq[Expression], names: Seq[Expression]) inputArrays.map(_.numElements()).max } - val result = new Array[InternalRow](biggestCardinality) + val result = new Array[AnyRef](biggestCardinality) val zippedArrs: Seq[(ArrayData, Int)] = inputArrays.zipWithIndex for (i <- 0 until biggestCardinality) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala index 7005d663a3f96..574a61cf9c903 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala @@ -1072,7 +1072,7 @@ case class JsonObjectKeys(child: Expression) extends UnaryExpression with Codege // skip all the children of inner object or array parser.skipChildren() } - new GenericArrayData(arrayBufferOfKeys.toArray) + new GenericArrayData(arrayBufferOfKeys.toArray[Any]) } override protected def withNewChildInternal(newChild: Expression): JsonObjectKeys = diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala index 5b06741a2f54e..31e65cf0abc95 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala @@ -255,7 +255,7 @@ case class XPathList(xml: Expression, path: Expression) extends XPathExtract { override def nullSafeEval(xml: Any, path: Any): Any = { val nodeList = xpathUtil.evalNodeList(xml.asInstanceOf[UTF8String].toString, pathString) if (nodeList ne null) { - val ret = new Array[UTF8String](nodeList.getLength) + val ret = new Array[AnyRef](nodeList.getLength) var i = 0 while (i < nodeList.getLength) { ret(i) = UTF8String.fromString(nodeList.item(i).getNodeValue) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala index 26de4cc7ad1c8..13129d44fe0c2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala @@ -215,7 +215,7 @@ class JacksonParser( ) } - Some(InternalRow(new GenericArrayData(res.toArray))) + Some(InternalRow(new GenericArrayData(res.toArray[Any]))) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/CommandUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/CommandUtils.scala index 5a9adf8ab553d..91454c79df600 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/CommandUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/CommandUtils.scala @@ -330,7 +330,7 @@ object CommandUtils extends Logging { val attributePercentiles = mutable.HashMap[Attribute, ArrayData]() if (attrsToGenHistogram.nonEmpty) { val percentiles = (0 to conf.histogramNumBins) - .map(i => i.toDouble / conf.histogramNumBins).toArray + .map(i => i.toDouble / conf.histogramNumBins).toArray[Any] val namedExprs = attrsToGenHistogram.map { attr => val aggFunc = From 3cb8d6e59999e5525374c62f964c57657935311c Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Thu, 12 Sep 2024 08:38:10 +0900 Subject: [PATCH 003/189] [SPARK-49548][CONNECT] Replace coarse-locking in SparkConnectSessionManager with ConcurrentMap ### What changes were proposed in this pull request? Replace the coarse-locking in SparkConnectSessionManager with ConcurrentMap in order to minimise lock contention when there are many sessions. ### Why are the changes needed? It is a spin-off from https://github.com/apache/spark/pull/48034 where https://github.com/apache/spark/pull/48034 addresses many-execution cases whereas this addresses many-session situations. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing test cases. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48036 from changgyoopark-db/SPARK-49548. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../service/SparkConnectSessionManager.scala | 99 +++++++++---------- 1 file changed, 49 insertions(+), 50 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala index edaaa640bf12e..fec01813de6e2 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.connect.service import java.util.UUID -import java.util.concurrent.{Executors, ScheduledExecutorService, TimeUnit} +import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit} import javax.annotation.concurrent.GuardedBy import scala.collection.mutable @@ -42,8 +42,8 @@ class SparkConnectSessionManager extends Logging { private val sessionsLock = new Object - @GuardedBy("sessionsLock") - private val sessionStore = mutable.HashMap[SessionKey, SessionHolder]() + private val sessionStore: ConcurrentMap[SessionKey, SessionHolder] = + new ConcurrentHashMap[SessionKey, SessionHolder]() private val closedSessionsCache = CacheBuilder @@ -52,6 +52,7 @@ class SparkConnectSessionManager extends Logging { .build[SessionKey, SessionHolderInfo]() /** Executor for the periodic maintenance */ + @GuardedBy("sessionsLock") private var scheduledExecutor: Option[ScheduledExecutorService] = None private def validateSessionId( @@ -121,43 +122,39 @@ class SparkConnectSessionManager extends Logging { private def getSession(key: SessionKey, default: Option[() => SessionHolder]): SessionHolder = { schedulePeriodicChecks() // Starts the maintenance thread if it hasn't started yet. - sessionsLock.synchronized { - // try to get existing session from store - val sessionOpt = sessionStore.get(key) - // create using default if missing - val session = sessionOpt match { - case Some(s) => s - case None => - default match { - case Some(callable) => - val session = callable() - sessionStore.put(key, session) - session - case None => - null - } - } - // record access time before returning - session match { - case null => - null - case s: SessionHolder => - s.updateAccessTime() - s - } + // Get the existing session from the store or create a new one. + val session = default match { + case Some(callable) => + sessionStore.computeIfAbsent(key, _ => callable()) + case None => + sessionStore.get(key) } + + // Record the access time before returning the session holder. + if (session != null) { + session.updateAccessTime() + } + + session } // Removes session from sessionStore and returns it. private def removeSessionHolder(key: SessionKey): Option[SessionHolder] = { var sessionHolder: Option[SessionHolder] = None - sessionsLock.synchronized { - sessionHolder = sessionStore.remove(key) - sessionHolder.foreach { s => - // Put into closedSessionsCache, so that it cannot get accidentally recreated - // by getOrCreateIsolatedSession. - closedSessionsCache.put(s.key, s.getSessionHolderInfo) - } + + // The session holder should remain in the session store until it is added to the closed session + // cache, because of a subtle data race: a new session with the same key can be created if the + // closed session cache does not contain the key right after the key has been removed from the + // session store. + sessionHolder = Option(sessionStore.get(key)) + + sessionHolder.foreach { s => + // Put into closedSessionsCache to prevent the same session from being recreated by + // getOrCreateIsolatedSession. + closedSessionsCache.put(s.key, s.getSessionHolderInfo) + + // Then, remove the session holder from the session store. + sessionStore.remove(key) } sessionHolder } @@ -176,21 +173,24 @@ class SparkConnectSessionManager extends Logging { sessionHolder.foreach(shutdownSessionHolder(_)) } - private[connect] def shutdown(): Unit = sessionsLock.synchronized { - scheduledExecutor.foreach { executor => - ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) + private[connect] def shutdown(): Unit = { + sessionsLock.synchronized { + scheduledExecutor.foreach { executor => + ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) + } + scheduledExecutor = None } - scheduledExecutor = None + // note: this does not cleanly shut down the sessions, but the server is shutting down. sessionStore.clear() closedSessionsCache.invalidateAll() } - def listActiveSessions: Seq[SessionHolderInfo] = sessionsLock.synchronized { - sessionStore.values.map(_.getSessionHolderInfo).toSeq + def listActiveSessions: Seq[SessionHolderInfo] = { + sessionStore.values().asScala.map(_.getSessionHolderInfo).toSeq } - def listClosedSessions: Seq[SessionHolderInfo] = sessionsLock.synchronized { + def listClosedSessions: Seq[SessionHolderInfo] = { closedSessionsCache.asMap.asScala.values.toSeq } @@ -246,18 +246,17 @@ class SparkConnectSessionManager extends Logging { timeoutMs != -1 && info.lastAccessTimeMs + timeoutMs <= nowMs } - sessionsLock.synchronized { - val nowMs = System.currentTimeMillis() - sessionStore.values.foreach { sessionHolder => - if (shouldExpire(sessionHolder.getSessionHolderInfo, nowMs)) { - toRemove += sessionHolder - } + val nowMs = System.currentTimeMillis() + sessionStore.forEach((_, sessionHolder) => { + if (shouldExpire(sessionHolder.getSessionHolderInfo, nowMs)) { + toRemove += sessionHolder } - } + }) + // .. and remove them. toRemove.foreach { sessionHolder => // This doesn't use closeSession to be able to do the extra last chance check under lock. - val removedSession = sessionsLock.synchronized { + val removedSession = { // Last chance - check expiration time and remove under lock if expired. val info = sessionHolder.getSessionHolderInfo if (shouldExpire(info, System.currentTimeMillis())) { @@ -309,7 +308,7 @@ class SparkConnectSessionManager extends Logging { /** * Used for testing */ - private[connect] def invalidateAllSessions(): Unit = sessionsLock.synchronized { + private[connect] def invalidateAllSessions(): Unit = { periodicMaintenance(defaultInactiveTimeoutMs = 0L, ignoreCustomTimeout = true) assert(sessionStore.isEmpty) closedSessionsCache.invalidateAll() From b466f32077e3f241cb8dfcd926098e4594635ace Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Thu, 12 Sep 2024 08:38:50 +0900 Subject: [PATCH 004/189] [SPARK-49544][CONNECT] Replace coarse-locking in SparkConnectExecutionManager with ConcurrentMap ### What changes were proposed in this pull request? Replace the coarse-locking mechanism implemented in SparkConnectExecutionManager with ConcurrentMap in order to ameliorate lock contention. ### Why are the changes needed? When there are too many threads, e.g., ~10K threads on a 4-core node, OS scheduling may cause priority inversion that leads to a serious performance problems, e.g., a 1000s delay when reattaching to an execute holder. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing test cases. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48034 from changgyoopark-db/SPARK-49544. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../sql/connect/service/ExecuteHolder.scala | 28 +-- .../SparkConnectExecutionManager.scala | 185 +++++++++++------- .../service/ExecuteEventsManagerSuite.scala | 3 +- 3 files changed, 123 insertions(+), 93 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala index ec7ebbe92d72e..dc349c3e33251 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala @@ -17,12 +17,10 @@ package org.apache.spark.sql.connect.service -import java.util.UUID - import scala.collection.mutable import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkEnv, SparkSQLException} +import org.apache.spark.SparkEnv import org.apache.spark.connect.proto import org.apache.spark.internal.Logging import org.apache.spark.sql.Observation @@ -35,30 +33,19 @@ import org.apache.spark.util.SystemClock * Object used to hold the Spark Connect execution state. */ private[connect] class ExecuteHolder( + val executeKey: ExecuteKey, val request: proto.ExecutePlanRequest, val sessionHolder: SessionHolder) extends Logging { val session = sessionHolder.session - val operationId = if (request.hasOperationId) { - try { - UUID.fromString(request.getOperationId).toString - } catch { - case _: IllegalArgumentException => - throw new SparkSQLException( - errorClass = "INVALID_HANDLE.FORMAT", - messageParameters = Map("handle" -> request.getOperationId)) - } - } else { - UUID.randomUUID().toString - } - /** * Tag that is set for this execution on SparkContext, via SparkContext.addJobTag. Used * (internally) for cancellation of the Spark Jobs ran by this execution. */ - val jobTag = ExecuteJobTag(sessionHolder.userId, sessionHolder.sessionId, operationId) + val jobTag = + ExecuteJobTag(sessionHolder.userId, sessionHolder.sessionId, executeKey.operationId) /** * Tags set by Spark Connect client users via SparkSession.addTag. Used to identify and group @@ -278,7 +265,7 @@ private[connect] class ExecuteHolder( request = request, userId = sessionHolder.userId, sessionId = sessionHolder.sessionId, - operationId = operationId, + operationId = executeKey.operationId, jobTag = jobTag, sparkSessionTags = sparkSessionTags, reattachable = reattachable, @@ -289,7 +276,10 @@ private[connect] class ExecuteHolder( } /** Get key used by SparkConnectExecutionManager global tracker. */ - def key: ExecuteKey = ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, operationId) + def key: ExecuteKey = executeKey + + /** Get the operation ID. */ + def operationId: String = key.operationId } /** Used to identify ExecuteHolder jobTag among SparkContext.SPARK_JOB_TAGS. */ diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala index 6681a5f509c6e..61b41f932199e 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala @@ -17,7 +17,8 @@ package org.apache.spark.sql.connect.service -import java.util.concurrent.{Executors, ScheduledExecutorService, TimeUnit} +import java.util.UUID +import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit} import javax.annotation.concurrent.GuardedBy import scala.collection.mutable @@ -36,6 +37,24 @@ import org.apache.spark.util.ThreadUtils // Unique key identifying execution by combination of user, session and operation id case class ExecuteKey(userId: String, sessionId: String, operationId: String) +object ExecuteKey { + def apply(request: proto.ExecutePlanRequest, sessionHolder: SessionHolder): ExecuteKey = { + val operationId = if (request.hasOperationId) { + try { + UUID.fromString(request.getOperationId).toString + } catch { + case _: IllegalArgumentException => + throw new SparkSQLException( + errorClass = "INVALID_HANDLE.FORMAT", + messageParameters = Map("handle" -> request.getOperationId)) + } + } else { + UUID.randomUUID().toString + } + ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, operationId) + } +} + /** * Global tracker of all ExecuteHolder executions. * @@ -44,10 +63,9 @@ case class ExecuteKey(userId: String, sessionId: String, operationId: String) */ private[connect] class SparkConnectExecutionManager() extends Logging { - /** Hash table containing all current executions. Guarded by executionsLock. */ - @GuardedBy("executionsLock") - private val executions: mutable.HashMap[ExecuteKey, ExecuteHolder] = - new mutable.HashMap[ExecuteKey, ExecuteHolder]() + /** Concurrent hash table containing all the current executions. */ + private val executions: ConcurrentMap[ExecuteKey, ExecuteHolder] = + new ConcurrentHashMap[ExecuteKey, ExecuteHolder]() private val executionsLock = new Object /** Graveyard of tombstones of executions that were abandoned and removed. */ @@ -61,6 +79,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { private var lastExecutionTimeMs: Option[Long] = Some(System.currentTimeMillis()) /** Executor for the periodic maintenance */ + @GuardedBy("executionsLock") private var scheduledExecutor: Option[ScheduledExecutorService] = None /** @@ -76,27 +95,35 @@ private[connect] class SparkConnectExecutionManager() extends Logging { request.getUserContext.getUserId, request.getSessionId, previousSessionId) - val executeHolder = new ExecuteHolder(request, sessionHolder) + val executeKey = ExecuteKey(request, sessionHolder) + val executeHolder = executions.compute( + executeKey, + (executeKey, oldExecuteHolder) => { + // Check if the operation already exists, either in the active execution map, or in the + // graveyard of tombstones of executions that have been abandoned. The latter is to prevent + // double executions when the client retries, thinking it never reached the server, but in + // fact it did, and already got removed as abandoned. + if (oldExecuteHolder != null) { + throw new SparkSQLException( + errorClass = "INVALID_HANDLE.OPERATION_ALREADY_EXISTS", + messageParameters = Map("handle" -> executeKey.operationId)) + } + if (getAbandonedTombstone(executeKey).isDefined) { + throw new SparkSQLException( + errorClass = "INVALID_HANDLE.OPERATION_ABANDONED", + messageParameters = Map("handle" -> executeKey.operationId)) + } + new ExecuteHolder(executeKey, request, sessionHolder) + }) + + sessionHolder.addExecuteHolder(executeHolder) + executionsLock.synchronized { - // Check if the operation already exists, both in active executions, and in the graveyard - // of tombstones of executions that have been abandoned. - // The latter is to prevent double execution when a client retries execution, thinking it - // never reached the server, but in fact it did, and already got removed as abandoned. - if (executions.get(executeHolder.key).isDefined) { - throw new SparkSQLException( - errorClass = "INVALID_HANDLE.OPERATION_ALREADY_EXISTS", - messageParameters = Map("handle" -> executeHolder.operationId)) - } - if (getAbandonedTombstone(executeHolder.key).isDefined) { - throw new SparkSQLException( - errorClass = "INVALID_HANDLE.OPERATION_ABANDONED", - messageParameters = Map("handle" -> executeHolder.operationId)) + if (!executions.isEmpty()) { + lastExecutionTimeMs = None } - sessionHolder.addExecuteHolder(executeHolder) - executions.put(executeHolder.key, executeHolder) - lastExecutionTimeMs = None - logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, executeHolder.key)} is created.") } + logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, executeHolder.key)} is created.") schedulePeriodicChecks() // Starts the maintenance thread if it hasn't started. @@ -108,43 +135,50 @@ private[connect] class SparkConnectExecutionManager() extends Logging { * execution if still running, free all resources. */ private[connect] def removeExecuteHolder(key: ExecuteKey, abandoned: Boolean = false): Unit = { - var executeHolder: Option[ExecuteHolder] = None + val executeHolder = executions.get(key) + + if (executeHolder == null) { + return + } + + // Put into abandonedTombstones before removing it from executions, so that the client ends up + // getting an INVALID_HANDLE.OPERATION_ABANDONED error on a retry. + if (abandoned) { + abandonedTombstones.put(key, executeHolder.getExecuteInfo) + } + + // Remove the execution from the map *after* putting it in abandonedTombstones. + executions.remove(key) + executeHolder.sessionHolder.removeExecuteHolder(executeHolder.operationId) + executionsLock.synchronized { - executeHolder = executions.remove(key) - executeHolder.foreach { e => - // Put into abandonedTombstones under lock, so that if it's accessed it will end up - // with INVALID_HANDLE.OPERATION_ABANDONED error. - if (abandoned) { - abandonedTombstones.put(key, e.getExecuteInfo) - } - e.sessionHolder.removeExecuteHolder(e.operationId) - } if (executions.isEmpty) { lastExecutionTimeMs = Some(System.currentTimeMillis()) } - logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, key)} is removed.") } - // close the execution outside the lock - executeHolder.foreach { e => - e.close() - if (abandoned) { - // Update in abandonedTombstones: above it wasn't yet updated with closedTime etc. - abandonedTombstones.put(key, e.getExecuteInfo) - } + + logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, key)} is removed.") + + executeHolder.close() + if (abandoned) { + // Update in abandonedTombstones: above it wasn't yet updated with closedTime etc. + abandonedTombstones.put(key, executeHolder.getExecuteInfo) } } private[connect] def getExecuteHolder(key: ExecuteKey): Option[ExecuteHolder] = { - executionsLock.synchronized { - executions.get(key) - } + Option(executions.get(key)) } private[connect] def removeAllExecutionsForSession(key: SessionKey): Unit = { - val sessionExecutionHolders = executionsLock.synchronized { - executions.filter(_._2.sessionHolder.key == key) - } - sessionExecutionHolders.foreach { case (_, executeHolder) => + var sessionExecutionHolders = mutable.ArrayBuffer[ExecuteHolder]() + executions.forEach((_, executeHolder) => { + if (executeHolder.sessionHolder.key == key) { + sessionExecutionHolders += executeHolder + } + }) + + sessionExecutionHolders.foreach { executeHolder => val info = executeHolder.getExecuteInfo logInfo( log"Execution ${MDC(LogKeys.EXECUTE_INFO, info)} removed in removeSessionExecutions.") @@ -161,11 +195,11 @@ private[connect] class SparkConnectExecutionManager() extends Logging { * If there are no executions, return Left with System.currentTimeMillis of last active * execution. Otherwise return Right with list of ExecuteInfo of all executions. */ - def listActiveExecutions: Either[Long, Seq[ExecuteInfo]] = executionsLock.synchronized { + def listActiveExecutions: Either[Long, Seq[ExecuteInfo]] = { if (executions.isEmpty) { Left(lastExecutionTimeMs.get) } else { - Right(executions.values.map(_.getExecuteInfo).toBuffer.toSeq) + Right(executions.values().asScala.map(_.getExecuteInfo).toBuffer.toSeq) } } @@ -177,16 +211,22 @@ private[connect] class SparkConnectExecutionManager() extends Logging { abandonedTombstones.asMap.asScala.values.toSeq } - private[connect] def shutdown(): Unit = executionsLock.synchronized { - scheduledExecutor.foreach { executor => - ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) + private[connect] def shutdown(): Unit = { + executionsLock.synchronized { + scheduledExecutor.foreach { executor => + ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) + } + scheduledExecutor = None } - scheduledExecutor = None + // note: this does not cleanly shut down the executions, but the server is shutting down. executions.clear() abandonedTombstones.invalidateAll() - if (lastExecutionTimeMs.isEmpty) { - lastExecutionTimeMs = Some(System.currentTimeMillis()) + + executionsLock.synchronized { + if (lastExecutionTimeMs.isEmpty) { + lastExecutionTimeMs = Some(System.currentTimeMillis()) + } } } @@ -225,19 +265,18 @@ private[connect] class SparkConnectExecutionManager() extends Logging { // Find any detached executions that expired and should be removed. val toRemove = new mutable.ArrayBuffer[ExecuteHolder]() - executionsLock.synchronized { - val nowMs = System.currentTimeMillis() + val nowMs = System.currentTimeMillis() - executions.values.foreach { executeHolder => - executeHolder.lastAttachedRpcTimeMs match { - case Some(detached) => - if (detached + timeout <= nowMs) { - toRemove += executeHolder - } - case _ => // execution is active - } + executions.forEach((_, executeHolder) => { + executeHolder.lastAttachedRpcTimeMs match { + case Some(detached) => + if (detached + timeout <= nowMs) { + toRemove += executeHolder + } + case _ => // execution is active } - } + }) + // .. and remove them. toRemove.foreach { executeHolder => val info = executeHolder.getExecuteInfo @@ -250,16 +289,16 @@ private[connect] class SparkConnectExecutionManager() extends Logging { } // For testing. - private[connect] def setAllRPCsDeadline(deadlineMs: Long) = executionsLock.synchronized { - executions.values.foreach(_.setGrpcResponseSendersDeadline(deadlineMs)) + private[connect] def setAllRPCsDeadline(deadlineMs: Long) = { + executions.values().asScala.foreach(_.setGrpcResponseSendersDeadline(deadlineMs)) } // For testing. - private[connect] def interruptAllRPCs() = executionsLock.synchronized { - executions.values.foreach(_.interruptGrpcResponseSenders()) + private[connect] def interruptAllRPCs() = { + executions.values().asScala.foreach(_.interruptGrpcResponseSenders()) } - private[connect] def listExecuteHolders: Seq[ExecuteHolder] = executionsLock.synchronized { - executions.values.toSeq + private[connect] def listExecuteHolders: Seq[ExecuteHolder] = { + executions.values().asScala.toSeq } } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/ExecuteEventsManagerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/ExecuteEventsManagerSuite.scala index dbe8420eab03d..a9843e261fff8 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/ExecuteEventsManagerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/ExecuteEventsManagerSuite.scala @@ -374,7 +374,8 @@ class ExecuteEventsManagerSuite .setClientType(DEFAULT_CLIENT_TYPE) .build() - val executeHolder = new ExecuteHolder(executePlanRequest, sessionHolder) + val executeKey = ExecuteKey(executePlanRequest, sessionHolder) + val executeHolder = new ExecuteHolder(executeKey, executePlanRequest, sessionHolder) val eventsManager = ExecuteEventsManager(executeHolder, DEFAULT_CLOCK) eventsManager.status_(executeStatus) From 0f4d289b7932c91186d2da66095ebb41b6cd58c0 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Thu, 12 Sep 2024 02:11:28 +0200 Subject: [PATCH 005/189] [SPARK-48906][SQL] Introduce `SHOW COLLATIONS LIKE ...` syntax to show all collations ### What changes were proposed in this pull request? The pr aims to introduce `SHOW COLLATIONS LIKE ...` syntax to `show all collations`. ### Why are the changes needed? End-users will be able to obtain `collations` currently supported by the spark through SQL. Other databases, such as `MySQL`, also have similar syntax, ref: https://dev.mysql.com/doc/refman/9.0/en/show-collation.html image postgresql: https://database.guide/how-to-return-a-list-of-available-collations-in-postgresql/ ### Does this PR introduce _any_ user-facing change? Yes, end-users will be able to obtain `collation` currently supported by the spark through commands similar to the following |name|provider|version|binaryEquality|binaryOrdering|lowercaseEquality| | --------- | ----------- | ----------- | ----------- | ----------- | ----------- | ``` spark-sql (default)> SHOW COLLATIONS; UTF8_BINARY spark 1.0 true true false UTF8_LCASE spark 1.0 false false true ff_Adlm icu 153.120.0.0 false false false ff_Adlm_CI icu 153.120.0.0 false false false ff_Adlm_AI icu 153.120.0.0 false false false ff_Adlm_CI_AI icu 153.120.0.0 false false false ... spark-sql (default)> SHOW COLLATIONS LIKE '*UTF8_BINARY*'; UTF8_BINARY spark 1.0 true true false Time taken: 0.043 seconds, Fetched 1 row(s) ``` image ### How was this patch tested? Add new UT. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47364 from panbingkun/show_collation_syntax. Authored-by: panbingkun Signed-off-by: Max Gekk --- .../sql/catalyst/util/CollationFactory.java | 143 +++++++++++++++++- docs/sql-ref-ansi-compliance.md | 1 + .../spark/sql/catalyst/parser/SqlBaseLexer.g4 | 1 + .../sql/catalyst/parser/SqlBaseParser.g4 | 2 + .../sql/catalyst/catalog/SessionCatalog.scala | 15 +- .../ansi-sql-2016-reserved-keywords.txt | 1 + .../spark/sql/execution/SparkSqlParser.scala | 12 ++ .../command/ShowCollationsCommand.scala | 62 ++++++++ .../sql-tests/results/ansi/keywords.sql.out | 2 + .../sql-tests/results/keywords.sql.out | 1 + .../org/apache/spark/sql/CollationSuite.scala | 42 +++++ .../ThriftServerWithSparkContextSuite.scala | 2 +- 12 files changed, 278 insertions(+), 6 deletions(-) create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java index 5640a2468d02e..4b88e15e8ed72 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java @@ -23,12 +23,14 @@ import java.util.function.Function; import java.util.function.BiFunction; import java.util.function.ToLongFunction; +import java.util.stream.Stream; +import com.ibm.icu.text.CollationKey; +import com.ibm.icu.text.Collator; import com.ibm.icu.text.RuleBasedCollator; import com.ibm.icu.text.StringSearch; import com.ibm.icu.util.ULocale; -import com.ibm.icu.text.CollationKey; -import com.ibm.icu.text.Collator; +import com.ibm.icu.util.VersionInfo; import org.apache.spark.SparkException; import org.apache.spark.unsafe.types.UTF8String; @@ -88,6 +90,17 @@ public Optional getVersion() { } } + public record CollationMeta( + String catalog, + String schema, + String collationName, + String language, + String country, + String icuVersion, + String padAttribute, + boolean accentSensitivity, + boolean caseSensitivity) { } + /** * Entry encapsulating all information about a collation. */ @@ -342,6 +355,23 @@ private static int collationNameToId(String collationName) throws SparkException } protected abstract Collation buildCollation(); + + protected abstract CollationMeta buildCollationMeta(); + + static List listCollations() { + return Stream.concat( + CollationSpecUTF8.listCollations().stream(), + CollationSpecICU.listCollations().stream()).toList(); + } + + static CollationMeta loadCollationMeta(CollationIdentifier collationIdentifier) { + CollationMeta collationSpecUTF8 = + CollationSpecUTF8.loadCollationMeta(collationIdentifier); + if (collationSpecUTF8 == null) { + return CollationSpecICU.loadCollationMeta(collationIdentifier); + } + return collationSpecUTF8; + } } private static class CollationSpecUTF8 extends CollationSpec { @@ -364,6 +394,9 @@ private enum CaseSensitivity { */ private static final int CASE_SENSITIVITY_MASK = 0b1; + private static final String UTF8_BINARY_COLLATION_NAME = "UTF8_BINARY"; + private static final String UTF8_LCASE_COLLATION_NAME = "UTF8_LCASE"; + private static final int UTF8_BINARY_COLLATION_ID = new CollationSpecUTF8(CaseSensitivity.UNSPECIFIED).collationId; private static final int UTF8_LCASE_COLLATION_ID = @@ -406,7 +439,7 @@ private static CollationSpecUTF8 fromCollationId(int collationId) { protected Collation buildCollation() { if (collationId == UTF8_BINARY_COLLATION_ID) { return new Collation( - "UTF8_BINARY", + UTF8_BINARY_COLLATION_NAME, PROVIDER_SPARK, null, UTF8String::binaryCompare, @@ -417,7 +450,7 @@ protected Collation buildCollation() { /* supportsLowercaseEquality = */ false); } else { return new Collation( - "UTF8_LCASE", + UTF8_LCASE_COLLATION_NAME, PROVIDER_SPARK, null, CollationAwareUTF8String::compareLowerCase, @@ -428,6 +461,52 @@ protected Collation buildCollation() { /* supportsLowercaseEquality = */ true); } } + + @Override + protected CollationMeta buildCollationMeta() { + if (collationId == UTF8_BINARY_COLLATION_ID) { + return new CollationMeta( + CATALOG, + SCHEMA, + UTF8_BINARY_COLLATION_NAME, + /* language = */ null, + /* country = */ null, + /* icuVersion = */ null, + COLLATION_PAD_ATTRIBUTE, + /* accentSensitivity = */ true, + /* caseSensitivity = */ true); + } else { + return new CollationMeta( + CATALOG, + SCHEMA, + UTF8_LCASE_COLLATION_NAME, + /* language = */ null, + /* country = */ null, + /* icuVersion = */ null, + COLLATION_PAD_ATTRIBUTE, + /* accentSensitivity = */ true, + /* caseSensitivity = */ false); + } + } + + static List listCollations() { + CollationIdentifier UTF8_BINARY_COLLATION_IDENT = + new CollationIdentifier(PROVIDER_SPARK, UTF8_BINARY_COLLATION_NAME, "1.0"); + CollationIdentifier UTF8_LCASE_COLLATION_IDENT = + new CollationIdentifier(PROVIDER_SPARK, UTF8_LCASE_COLLATION_NAME, "1.0"); + return Arrays.asList(UTF8_BINARY_COLLATION_IDENT, UTF8_LCASE_COLLATION_IDENT); + } + + static CollationMeta loadCollationMeta(CollationIdentifier collationIdentifier) { + try { + int collationId = CollationSpecUTF8.collationNameToId( + collationIdentifier.name, collationIdentifier.name.toUpperCase()); + return CollationSpecUTF8.fromCollationId(collationId).buildCollationMeta(); + } catch (SparkException ignored) { + // ignore + return null; + } + } } private static class CollationSpecICU extends CollationSpec { @@ -684,6 +763,20 @@ protected Collation buildCollation() { /* supportsLowercaseEquality = */ false); } + @Override + protected CollationMeta buildCollationMeta() { + return new CollationMeta( + CATALOG, + SCHEMA, + collationName(), + ICULocaleMap.get(locale).getDisplayLanguage(), + ICULocaleMap.get(locale).getDisplayCountry(), + VersionInfo.ICU_VERSION.toString(), + COLLATION_PAD_ATTRIBUTE, + caseSensitivity == CaseSensitivity.CS, + accentSensitivity == AccentSensitivity.AS); + } + /** * Compute normalized collation name. Components of collation name are given in order: * - Locale name @@ -704,6 +797,37 @@ private String collationName() { } return builder.toString(); } + + private static List allCollationNames() { + List collationNames = new ArrayList<>(); + for (String locale: ICULocaleToId.keySet()) { + // CaseSensitivity.CS + AccentSensitivity.AS + collationNames.add(locale); + // CaseSensitivity.CS + AccentSensitivity.AI + collationNames.add(locale + "_AI"); + // CaseSensitivity.CI + AccentSensitivity.AS + collationNames.add(locale + "_CI"); + // CaseSensitivity.CI + AccentSensitivity.AI + collationNames.add(locale + "_CI_AI"); + } + return collationNames.stream().sorted().toList(); + } + + static List listCollations() { + return allCollationNames().stream().map(name -> + new CollationIdentifier(PROVIDER_ICU, name, VersionInfo.ICU_VERSION.toString())).toList(); + } + + static CollationMeta loadCollationMeta(CollationIdentifier collationIdentifier) { + try { + int collationId = CollationSpecICU.collationNameToId( + collationIdentifier.name, collationIdentifier.name.toUpperCase()); + return CollationSpecICU.fromCollationId(collationId).buildCollationMeta(); + } catch (SparkException ignored) { + // ignore + return null; + } + } } /** @@ -730,9 +854,12 @@ public CollationIdentifier identifier() { } } + public static final String CATALOG = "SYSTEM"; + public static final String SCHEMA = "BUILTIN"; public static final String PROVIDER_SPARK = "spark"; public static final String PROVIDER_ICU = "icu"; public static final List SUPPORTED_PROVIDERS = List.of(PROVIDER_SPARK, PROVIDER_ICU); + public static final String COLLATION_PAD_ATTRIBUTE = "NO_PAD"; public static final int UTF8_BINARY_COLLATION_ID = Collation.CollationSpecUTF8.UTF8_BINARY_COLLATION_ID; @@ -923,4 +1050,12 @@ public static String getClosestSuggestionsOnInvalidName( return String.join(", ", suggestions); } + + public static List listCollations() { + return Collation.CollationSpec.listCollations(); + } + + public static CollationMeta loadCollationMeta(CollationIdentifier collationIdentifier) { + return Collation.CollationSpec.loadCollationMeta(collationIdentifier); + } } diff --git a/docs/sql-ref-ansi-compliance.md b/docs/sql-ref-ansi-compliance.md index 3fa67036fd04b..fe5ddf27bf6c4 100644 --- a/docs/sql-ref-ansi-compliance.md +++ b/docs/sql-ref-ansi-compliance.md @@ -442,6 +442,7 @@ Below is a list of all the keywords in Spark SQL. |CODEGEN|non-reserved|non-reserved|non-reserved| |COLLATE|reserved|non-reserved|reserved| |COLLATION|reserved|non-reserved|reserved| +|COLLATIONS|reserved|non-reserved|reserved| |COLLECTION|non-reserved|non-reserved|non-reserved| |COLUMN|reserved|non-reserved|reserved| |COLUMNS|non-reserved|non-reserved|non-reserved| diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index 28ebaeaaed6d0..9ea213f3bf4a6 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -162,6 +162,7 @@ CLUSTERED: 'CLUSTERED'; CODEGEN: 'CODEGEN'; COLLATE: 'COLLATE'; COLLATION: 'COLLATION'; +COLLATIONS: 'COLLATIONS'; COLLECTION: 'COLLECTION'; COLUMN: 'COLUMN'; COLUMNS: 'COLUMNS'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index e9fc6c3ca4f2e..42f0094de3515 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -260,6 +260,7 @@ statement | SHOW PARTITIONS identifierReference partitionSpec? #showPartitions | SHOW identifier? FUNCTIONS ((FROM | IN) ns=identifierReference)? (LIKE? (legacy=multipartIdentifier | pattern=stringLit))? #showFunctions + | SHOW COLLATIONS (LIKE? pattern=stringLit)? #showCollations | SHOW CREATE TABLE identifierReference (AS SERDE)? #showCreateTable | SHOW CURRENT namespace #showCurrentNamespace | SHOW CATALOGS (LIKE? pattern=stringLit)? #showCatalogs @@ -1837,6 +1838,7 @@ nonReserved | CODEGEN | COLLATE | COLLATION + | COLLATIONS | COLLECTION | COLUMN | COLUMNS diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala index d3a6cb6ae2845..5c14e261fafc8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala @@ -24,6 +24,7 @@ import java.util.concurrent.TimeUnit import javax.annotation.concurrent.GuardedBy import scala.collection.mutable +import scala.jdk.CollectionConverters.CollectionHasAsScala import scala.util.{Failure, Success, Try} import com.google.common.cache.{Cache, CacheBuilder} @@ -39,7 +40,8 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Expression, Expre import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParserInterface} import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, SubqueryAlias, View} import org.apache.spark.sql.catalyst.trees.CurrentOrigin -import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, StringUtils} +import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, CollationFactory, StringUtils} +import org.apache.spark.sql.catalyst.util.CollationFactory.CollationMeta import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} @@ -1899,6 +1901,17 @@ class SessionCatalog( .filter(isTemporaryFunction) } + /** + * List all built-in collations with the given pattern. + */ + def listCollations(pattern: Option[String]): Seq[CollationMeta] = { + val collationIdentifiers = CollationFactory.listCollations().asScala.toSeq + val filteredCollationNames = StringUtils.filterPattern( + collationIdentifiers.map(_.getName), pattern.getOrElse("*")).toSet + collationIdentifiers.filter(ident => filteredCollationNames.contains(ident.getName)).map( + CollationFactory.loadCollationMeta) + } + // ----------------- // | Other methods | // ----------------- diff --git a/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt b/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt index 46da60b7897b8..452cf930525bc 100644 --- a/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt +++ b/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt @@ -48,6 +48,7 @@ CLOSE COALESCE COLLATE COLLATION +COLLATIONS COLLECT COLUMN COMMIT diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala index a8261e5d98ba0..640abaea58abe 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala @@ -1096,4 +1096,16 @@ class SparkSqlAstBuilder extends AstBuilder { withIdentClause(ctx.identifierReference(), UnresolvedNamespace(_)), cleanedProperties) } + + /** + * Create a [[ShowCollationsCommand]] command. + * Expected format: + * {{{ + * SHOW COLLATIONS (LIKE? pattern=stringLit)?; + * }}} + */ + override def visitShowCollations(ctx: ShowCollationsContext): LogicalPlan = withOrigin(ctx) { + val pattern = Option(ctx.pattern).map(x => string(visitStringLit(x))) + ShowCollationsCommand(pattern) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala new file mode 100644 index 0000000000000..179a841b013bd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.command + +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.catalyst.util.CollationFactory.CollationMeta +import org.apache.spark.sql.types.StringType + +/** + * A command for `SHOW COLLATIONS`. + * + * The syntax of this command is: + * {{{ + * SHOW COLLATIONS (LIKE? pattern=stringLit)?; + * }}} + */ +case class ShowCollationsCommand(pattern: Option[String]) extends LeafRunnableCommand { + + override val output: Seq[Attribute] = Seq( + AttributeReference("COLLATION_CATALOG", StringType, nullable = false)(), + AttributeReference("COLLATION_SCHEMA", StringType, nullable = false)(), + AttributeReference("COLLATION_NAME", StringType, nullable = false)(), + AttributeReference("LANGUAGE", StringType)(), + AttributeReference("COUNTRY", StringType)(), + AttributeReference("ACCENT_SENSITIVITY", StringType, nullable = false)(), + AttributeReference("CASE_SENSITIVITY", StringType, nullable = false)(), + AttributeReference("PAD_ATTRIBUTE", StringType, nullable = false)(), + AttributeReference("ICU_VERSION", StringType)()) + + override def run(sparkSession: SparkSession): Seq[Row] = { + val systemCollations: Seq[CollationMeta] = + sparkSession.sessionState.catalog.listCollations(pattern) + + systemCollations.map(m => Row( + m.catalog, + m.schema, + m.collationName, + m.language, + m.country, + if (m.accentSensitivity) "ACCENT_SENSITIVE" else "ACCENT_INSENSITIVE", + if (m.caseSensitivity) "CASE_SENSITIVE" else "CASE_INSENSITIVE", + m.padAttribute, + m.icuVersion + )) + } +} diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out index e6a36ac2445cf..81ccc0f9efc13 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out @@ -48,6 +48,7 @@ CLUSTERED false CODEGEN false COLLATE true COLLATION true +COLLATIONS true COLLECTION false COLUMN true COLUMNS false @@ -381,6 +382,7 @@ CAST CHECK COLLATE COLLATION +COLLATIONS COLUMN CONSTRAINT CREATE diff --git a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out index 19816c8252c91..e145c57332eb2 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out @@ -48,6 +48,7 @@ CLUSTERED false CODEGEN false COLLATE false COLLATION false +COLLATIONS false COLLECTION false COLUMN false COLUMNS false diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index a61be9eca8c31..b25cddb80762a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -1624,4 +1624,46 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { } } } + + test("show collations") { + assert(sql("SHOW COLLATIONS").collect().length >= 562) + + // verify that the output ordering is as expected (UTF8_BINARY, UTF8_LCASE, etc.) + val df = sql("SHOW COLLATIONS").limit(10) + checkAnswer(df, + Seq(Row("SYSTEM", "BUILTIN", "UTF8_BINARY", null, null, + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", null), + Row("SYSTEM", "BUILTIN", "UTF8_LCASE", null, null, + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", null), + Row("SYSTEM", "BUILTIN", "UNICODE", "", "", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_AI", "", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_CI", "", "", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_CI_AI", "", "", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af", "Afrikaans", "", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_AI", "Afrikaans", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_CI", "Afrikaans", "", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_CI_AI", "Afrikaans", "", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + + checkAnswer(sql("SHOW COLLATIONS LIKE '*UTF8_BINARY*'"), + Row("SYSTEM", "BUILTIN", "UTF8_BINARY", null, null, + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", null)) + + checkAnswer(sql("SHOW COLLATIONS '*zh_Hant_HKG*'"), + Seq(Row("SYSTEM", "BUILTIN", "zh_Hant_HKG", "Chinese", "Hong Kong SAR China", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_AI", "Chinese", "Hong Kong SAR China", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI", "Chinese", "Hong Kong SAR China", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI_AI", "Chinese", "Hong Kong SAR China", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + } } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala index 6f0b6bccac309..edef6371be8ae 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala @@ -214,7 +214,7 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer { val sessionHandle = client.openSession(user, "") val infoValue = client.getInfo(sessionHandle, GetInfoType.CLI_ODBC_KEYWORDS) // scalastyle:off line.size.limit - assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") + assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") // scalastyle:on line.size.limit } } From 8023504e69fdd037dea002e961b960fd9fa662ba Mon Sep 17 00:00:00 2001 From: Eric Marnadi Date: Thu, 12 Sep 2024 12:01:08 +0900 Subject: [PATCH 006/189] [SPARK-49594][SS] Adding check on whether columnFamilies were added or removed to write StateSchemaV3 file ### What changes were proposed in this pull request? Up until this [PR](https://github.com/apache/spark/pull/47880) that enabled deleteIfExists, we changed the condition on which we throw an error. However, in doing so, we are not writing schema files whenever we add or remove column families, which is functionally incorrect. Additionally, we were initially always writing the newSchemaFilePath to the OperatorStateMetadata upon every new query run, when we should only do this if the schema changes. ### Why are the changes needed? These changes are needed because we want to write a schema file out every time we add or remove column families. Also, we want to make sure that we point to the old schema file for the current metadata file if the schema has not changed between this run and the last one, as opposed to populating the metadata with a new schema file path every time, even if this file is not created. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Amended unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #48067 from ericm-db/add-remove-cf. Authored-by: Eric Marnadi Signed-off-by: Jungtaek Lim --- .../StateSchemaCompatibilityChecker.scala | 40 +++- .../streaming/TransformWithStateSuite.scala | 219 +++++++++++++++++- 2 files changed, 250 insertions(+), 9 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala index 90eb634689b23..3a1793f71794f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala @@ -27,6 +27,7 @@ import org.apache.spark.internal.{Logging, LogKeys, MDC} import org.apache.spark.sql.catalyst.util.UnsafeRowUtils import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, StatefulOperatorStateInfo} import org.apache.spark.sql.execution.streaming.state.SchemaHelper.{SchemaReader, SchemaWriter} +import org.apache.spark.sql.execution.streaming.state.StateSchemaCompatibilityChecker.SCHEMA_FORMAT_V3 import org.apache.spark.sql.internal.SessionState import org.apache.spark.sql.types.{DataType, StructType} @@ -95,7 +96,7 @@ class StateSchemaCompatibilityChecker( stateStoreColFamilySchema: List[StateStoreColFamilySchema], stateSchemaVersion: Int): Unit = { // Ensure that schema file path is passed explicitly for schema version 3 - if (stateSchemaVersion == 3 && newSchemaFilePath.isEmpty) { + if (stateSchemaVersion == SCHEMA_FORMAT_V3 && newSchemaFilePath.isEmpty) { throw new IllegalStateException("Schema file path is required for schema version 3") } @@ -186,8 +187,13 @@ class StateSchemaCompatibilityChecker( check(existingStateSchema, newSchema, ignoreValueSchema) } } + val colFamiliesAddedOrRemoved = + newStateSchemaList.map(_.colFamilyName) != existingStateSchemaList.map(_.colFamilyName) + if (stateSchemaVersion == SCHEMA_FORMAT_V3 && colFamiliesAddedOrRemoved) { + createSchemaFile(newStateSchemaList, stateSchemaVersion) + } // TODO: [SPARK-49535] Write Schema files after schema has changed for StateSchemaV3 - false + colFamiliesAddedOrRemoved } } @@ -196,6 +202,9 @@ class StateSchemaCompatibilityChecker( } object StateSchemaCompatibilityChecker { + + val SCHEMA_FORMAT_V3: Int = 3 + private def disallowBinaryInequalityColumn(schema: StructType): Unit = { if (!UnsafeRowUtils.isBinaryStable(schema)) { throw new SparkUnsupportedOperationException( @@ -275,10 +284,31 @@ object StateSchemaCompatibilityChecker { if (storeConf.stateSchemaCheckEnabled && result.isDefined) { throw result.get } - val schemaFileLocation = newSchemaFilePath match { - case Some(path) => path.toString - case None => checker.schemaFileLocation.toString + val schemaFileLocation = if (evolvedSchema) { + // if we are using the state schema v3, and we have + // evolved schema, this newSchemaFilePath should be defined + // and we want to populate the metadata with this file + if (stateSchemaVersion == SCHEMA_FORMAT_V3) { + newSchemaFilePath.get.toString + } else { + // if we are using any version less than v3, we have written + // the schema to this static location, which we will return + checker.schemaFileLocation.toString + } + } else { + // if we have not evolved schema (there has been a previous schema) + // and we are using state schema v3, this file path would be defined + // so we would just populate the next run's metadata file with this + // file path + if (stateSchemaVersion == SCHEMA_FORMAT_V3) { + oldSchemaFilePath.get.toString + } else { + // if we are using any version less than v3, we have written + // the schema to this static location, which we will return + checker.schemaFileLocation.toString + } } + StateSchemaValidationResult(evolvedSchema, schemaFileLocation) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala index a17f3847323d5..d0e255bb30499 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala @@ -1448,6 +1448,10 @@ class TransformWithStateSuite extends StateStoreMetricsTest TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString, SQLConf.MIN_BATCHES_TO_RETAIN.key -> "1") { withTempDir { chkptDir => + val stateOpIdPath = new Path(new Path(chkptDir.getCanonicalPath, "state"), "0") + val stateSchemaPath = getStateSchemaPath(stateOpIdPath) + + val metadataPath = OperatorStateMetadataV2.metadataDirPath(stateOpIdPath) // in this test case, we are changing the state spec back and forth // to trigger the writing of the schema and metadata files val inputData = MemoryStream[(String, String)] @@ -1483,6 +1487,11 @@ class TransformWithStateSuite extends StateStoreMetricsTest }, StopStream ) + // assert that a metadata and schema file has been written for each run + // as state variables have been deleted + assert(getFiles(metadataPath).length == 2) + assert(getFiles(stateSchemaPath).length == 2) + val result3 = inputData.toDS() .groupByKey(x => x._1) .transformWithState(new RunningCountMostRecentStatefulProcessor(), @@ -1512,10 +1521,6 @@ class TransformWithStateSuite extends StateStoreMetricsTest }, StopStream ) - val stateOpIdPath = new Path(new Path(chkptDir.getCanonicalPath, "state"), "0") - val stateSchemaPath = getStateSchemaPath(stateOpIdPath) - - val metadataPath = OperatorStateMetadataV2.metadataDirPath(stateOpIdPath) // by the end of the test, there have been 4 batches, // so the metadata and schema logs, and commitLog has been purged // for batches 0 and 1 so metadata and schema files exist for batches 0, 1, 2, 3 @@ -1527,6 +1532,116 @@ class TransformWithStateSuite extends StateStoreMetricsTest } } + test("transformWithState - verify that schema file is kept after metadata is purged") { + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString, + SQLConf.MIN_BATCHES_TO_RETAIN.key -> "2") { + withTempDir { chkptDir => + val stateOpIdPath = new Path(new Path(chkptDir.getCanonicalPath, "state"), "0") + val stateSchemaPath = getStateSchemaPath(stateOpIdPath) + + val metadataPath = OperatorStateMetadataV2.metadataDirPath(stateOpIdPath) + // in this test case, we are changing the state spec back and forth + // to trigger the writing of the schema and metadata files + val inputData = MemoryStream[(String, String)] + val result1 = inputData.toDS() + .groupByKey(x => x._1) + .transformWithState(new RunningCountMostRecentStatefulProcessor(), + TimeMode.None(), + OutputMode.Update()) + testStream(result1, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("a", "str1")), + CheckNewAnswer(("a", "1", "")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + testStream(result1, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("a", "str1")), + CheckNewAnswer(("a", "2", "str1")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + val result2 = inputData.toDS() + .groupByKey(x => x._1) + .transformWithState(new MostRecentStatefulProcessorWithDeletion(), + TimeMode.None(), + OutputMode.Update()) + testStream(result2, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("a", "str2")), + CheckNewAnswer(("a", "str1")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + assert(getFiles(metadataPath).length == 3) + assert(getFiles(stateSchemaPath).length == 2) + + val result3 = inputData.toDS() + .groupByKey(x => x._1) + .transformWithState(new RunningCountMostRecentStatefulProcessor(), + TimeMode.None(), + OutputMode.Update()) + testStream(result3, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("a", "str3")), + CheckNewAnswer(("a", "1", "str2")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + // metadata files should be kept for batches 1, 2, 3 + // schema files should be kept for batches 0, 2, 3 + assert(getFiles(metadataPath).length == 3) + assert(getFiles(stateSchemaPath).length == 3) + // we want to ensure that we can read batch 1 even though the + // metadata file for batch 0 was removed + val batch1Df = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "countState") + .option(StateSourceOptions.BATCH_ID, 1) + .load() + + val batch1AnsDf = batch1Df.selectExpr( + "key.value AS groupingKey", + "single_value.value AS valueId") + + checkAnswer(batch1AnsDf, Seq(Row("a", 2L))) + + val batch3Df = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, chkptDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "countState") + .option(StateSourceOptions.BATCH_ID, 3) + .load() + + val batch3AnsDf = batch3Df.selectExpr( + "key.value AS groupingKey", + "single_value.value AS valueId") + checkAnswer(batch3AnsDf, Seq(Row("a", 1L))) + } + } + } + test("state data source integration - value state supports time travel") { withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName, @@ -1708,6 +1823,102 @@ class TransformWithStateSuite extends StateStoreMetricsTest } } } + + test("transformWithState - verify that no metadata and schema logs are purged after" + + " removing column family") { + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString, + SQLConf.MIN_BATCHES_TO_RETAIN.key -> "3") { + withTempDir { chkptDir => + val inputData = MemoryStream[(String, String)] + val result1 = inputData.toDS() + .groupByKey(x => x._1) + .transformWithState(new RunningCountMostRecentStatefulProcessor(), + TimeMode.None(), + OutputMode.Update()) + testStream(result1, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("a", "str1")), + CheckNewAnswer(("a", "1", "")), + AddData(inputData, ("a", "str1")), + CheckNewAnswer(("a", "2", "str1")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + testStream(result1, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "1", "")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "2", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "3", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "4", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "5", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "6", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "7", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "8", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "9", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "10", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "11", "str1")), + AddData(inputData, ("b", "str1")), + CheckNewAnswer(("b", "12", "str1")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + val result2 = inputData.toDS() + .groupByKey(x => x._1) + .transformWithState(new MostRecentStatefulProcessorWithDeletion(), + TimeMode.None(), + OutputMode.Update()) + + testStream(result2, OutputMode.Update())( + StartStream(checkpointLocation = chkptDir.getCanonicalPath), + AddData(inputData, ("b", "str2")), + CheckNewAnswer(("b", "str1")), + AddData(inputData, ("b", "str3")), + CheckNewAnswer(("b", "str2")), + Execute { q => + eventually(timeout(Span(5, Seconds))) { + q.asInstanceOf[MicroBatchExecution].arePendingAsyncPurge should be(false) + } + }, + StopStream + ) + + val stateOpIdPath = new Path(new Path(chkptDir.getCanonicalPath, "state"), "0") + val stateSchemaPath = getStateSchemaPath(stateOpIdPath) + + val metadataPath = OperatorStateMetadataV2.metadataDirPath(stateOpIdPath) + + // Metadata files are written for batches 0, 2, and 14. + // Schema files are written for 0, 14 + // At the beginning of the last query run, the thresholdBatchId is 11. + // However, we would need both schema files to be preserved, if we want to + // be able to read from batch 11 onwards. + assert(getFiles(metadataPath).length == 2) + assert(getFiles(stateSchemaPath).length == 2) + } + } + } } class TransformWithStateValidationSuite extends StateStoreMetricsTest { From 19aad9ee36edad0906b8223074351bfb76237c0a Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 12 Sep 2024 07:17:28 -0700 Subject: [PATCH 007/189] [SPARK-49578][SQL][TESTS][FOLLOWUP] Regenerate Java 21 golden file for `postgreSQL/float4.sql` and `postgreSQL/int8.sql` ### What changes were proposed in this pull request? This pr regenerate Java 21 golden file for `postgreSQL/float4.sql` and `postgreSQL/int8.sql` to fix Java 21 daily test. ### Why are the changes needed? Fix Java 21 daily test: - https://github.com/apache/spark/actions/runs/10823897095/job/30030200710 ``` [info] - postgreSQL/float4.sql *** FAILED *** (1 second, 100 milliseconds) [info] postgreSQL/float4.sql [info] Expected "...arameters" : { [info] "[ansiConfig" : "\"spark.sql.ansi.enabled\"", [info] "]expression" : "'N A ...", but got "...arameters" : { [info] "[]expression" : "'N A ..." Result did not match for query #11 [info] SELECT float('N A N') (SQLQueryTestSuite.scala:663) ... [info] - postgreSQL/int8.sql *** FAILED *** (2 seconds, 474 milliseconds) [info] postgreSQL/int8.sql [info] Expected "...arameters" : { [info] "[ansiConfig" : "\"spark.sql.ansi.enabled\"", [info] "]sourceType" : "\"BIG...", but got "...arameters" : { [info] "[]sourceType" : "\"BIG..." Result did not match for query #66 [info] SELECT CAST(q1 AS int) FROM int8_tbl WHERE q2 <> 456 (SQLQueryTestSuite.scala:663) ... [info] *** 2 TESTS FAILED *** [error] Failed: Total 3559, Failed 2, Errors 0, Passed 3557, Ignored 4 [error] Failed tests: [error] org.apache.spark.sql.SQLQueryTestSuite [error] (sql / Test / test) sbt.TestsFailedException: Tests unsuccessful ``` ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? - Pass Github Acitons - Manual checked: `build/sbt "sql/testOnly org.apache.spark.sql.SQLQueryTestSuite" with Java 21, all test passed ` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48089 from LuciferYang/SPARK-49578-FOLLOWUP. Authored-by: yangjie01 Signed-off-by: Dongjoon Hyun --- .../sql-tests/results/postgreSQL/float4.sql.out.java21 | 7 ------- .../sql-tests/results/postgreSQL/int8.sql.out.java21 | 4 ---- 2 files changed, 11 deletions(-) diff --git a/sql/core/src/test/resources/sql-tests/results/postgreSQL/float4.sql.out.java21 b/sql/core/src/test/resources/sql-tests/results/postgreSQL/float4.sql.out.java21 index 6126411071bc1..3c2189c399639 100644 --- a/sql/core/src/test/resources/sql-tests/results/postgreSQL/float4.sql.out.java21 +++ b/sql/core/src/test/resources/sql-tests/results/postgreSQL/float4.sql.out.java21 @@ -97,7 +97,6 @@ org.apache.spark.SparkNumberFormatException "errorClass" : "CAST_INVALID_INPUT", "sqlState" : "22018", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "expression" : "'N A N'", "sourceType" : "\"STRING\"", "targetType" : "\"FLOAT\"" @@ -122,7 +121,6 @@ org.apache.spark.SparkNumberFormatException "errorClass" : "CAST_INVALID_INPUT", "sqlState" : "22018", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "expression" : "'NaN x'", "sourceType" : "\"STRING\"", "targetType" : "\"FLOAT\"" @@ -147,7 +145,6 @@ org.apache.spark.SparkNumberFormatException "errorClass" : "CAST_INVALID_INPUT", "sqlState" : "22018", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "expression" : "' INFINITY x'", "sourceType" : "\"STRING\"", "targetType" : "\"FLOAT\"" @@ -196,7 +193,6 @@ org.apache.spark.SparkNumberFormatException "errorClass" : "CAST_INVALID_INPUT", "sqlState" : "22018", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "expression" : "'nan'", "sourceType" : "\"STRING\"", "targetType" : "\"DECIMAL(10,0)\"" @@ -393,7 +389,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"FLOAT\"", "targetType" : "\"INT\"", "value" : "2.1474836E9" @@ -419,7 +414,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"FLOAT\"", "targetType" : "\"INT\"", "value" : "-2.147484E9" @@ -461,7 +455,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"FLOAT\"", "targetType" : "\"BIGINT\"", "value" : "-9.22338E18" diff --git a/sql/core/src/test/resources/sql-tests/results/postgreSQL/int8.sql.out.java21 b/sql/core/src/test/resources/sql-tests/results/postgreSQL/int8.sql.out.java21 index ee3f8625da8a4..e7df03dc8cadd 100755 --- a/sql/core/src/test/resources/sql-tests/results/postgreSQL/int8.sql.out.java21 +++ b/sql/core/src/test/resources/sql-tests/results/postgreSQL/int8.sql.out.java21 @@ -737,7 +737,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"BIGINT\"", "targetType" : "\"INT\"", "value" : "4567890123456789L" @@ -763,7 +762,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"BIGINT\"", "targetType" : "\"SMALLINT\"", "value" : "4567890123456789L" @@ -809,7 +807,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"DOUBLE\"", "targetType" : "\"BIGINT\"", "value" : "9.223372036854776E20D" @@ -898,7 +895,6 @@ org.apache.spark.SparkArithmeticException "errorClass" : "CAST_OVERFLOW", "sqlState" : "22003", "messageParameters" : { - "ansiConfig" : "\"spark.sql.ansi.enabled\"", "sourceType" : "\"BIGINT\"", "targetType" : "\"INT\"", "value" : "-9223372036854775808L" From c5c880e690c38b2bb597b7a38f20b32e2e2d272c Mon Sep 17 00:00:00 2001 From: cashmand Date: Thu, 12 Sep 2024 22:35:57 +0800 Subject: [PATCH 008/189] [SPARK-49591][SQL] Add Logical Type column to variant readme ### What changes were proposed in this pull request? Add a concept of logical type to the variant README.md, distinct from the physical encoding of a value. In particular, decimal and integer values are considered to be members of a single "Exact Numeric" type. ### Why are the changes needed? This is intended to describe and justify the existing Spark behaviour for Variant (e.g. stripping trailing zeros for decimal to string casts), not change it. (Although the SchemaOfVariant expression does not strictly follow this right now for numeric types, and should be updated to match it.) The motivation for introducing a single numeric type that encompasses integer and decimal values is to allow more flexibility in storage (particularly once shredding is introduced), and provide a simpler user surface, since there is not much benefit to distinguishing numeric precision/scale on a per-value basis. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? It is a documentation change. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48064 from cashmand/cashmand/SPARK-49591. Authored-by: cashmand Signed-off-by: Wenchen Fan --- common/variant/README.md | 44 +++++++++++++++++++++------------------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/common/variant/README.md b/common/variant/README.md index a66d708da75bf..4ed7c16f5b6ed 100644 --- a/common/variant/README.md +++ b/common/variant/README.md @@ -333,27 +333,27 @@ The Decimal type contains a scale, but no precision. The implied precision of a | Object | `2` | A collection of (string-key, variant-value) pairs | | Array | `3` | An ordered sequence of variant values | -| Primitive Type | Type ID | Equivalent Parquet Type | Binary format | -|-----------------------------|---------|-----------------------------|---------------------------------------------------------------------------------------------------------------------| -| null | `0` | any | none | -| boolean (True) | `1` | BOOLEAN | none | -| boolean (False) | `2` | BOOLEAN | none | -| int8 | `3` | INT(8, signed) | 1 byte | -| int16 | `4` | INT(16, signed) | 2 byte little-endian | -| int32 | `5` | INT(32, signed) | 4 byte little-endian | -| int64 | `6` | INT(64, signed) | 8 byte little-endian | -| double | `7` | DOUBLE | IEEE little-endian | -| decimal4 | `8` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | -| decimal8 | `9` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | -| decimal16 | `10` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | -| date | `11` | DATE | 4 byte little-endian | -| timestamp | `12` | TIMESTAMP(true, MICROS) | 8-byte little-endian | -| timestamp without time zone | `13` | TIMESTAMP(false, MICROS) | 8-byte little-endian | -| float | `14` | FLOAT | IEEE little-endian | -| binary | `15` | BINARY | 4 byte little-endian size, followed by bytes | -| string | `16` | STRING | 4 byte little-endian size, followed by UTF-8 encoded bytes | -| year-month interval | `19` | INT(32, signed)1 | 1 byte denoting start field (1 bit) and end field (1 bit) starting at LSB followed by 4-byte little-endian value. | -| day-time interval | `20` | INT(64, signed)1 | 1 byte denoting start field (2 bits) and end field (2 bits) starting at LSB followed by 8-byte little-endian value. | +| Logical Type | Physical Type | Type ID | Equivalent Parquet Type | Binary format | +|----------------------|-----------------------------|---------|-----------------------------|---------------------------------------------------------------------------------------------------------------------| +| NullType | null | `0` | any | none | +| Boolean | boolean (True) | `1` | BOOLEAN | none | +| Boolean | boolean (False) | `2` | BOOLEAN | none | +| Exact Numeric | int8 | `3` | INT(8, signed) | 1 byte | +| Exact Numeric | int16 | `4` | INT(16, signed) | 2 byte little-endian | +| Exact Numeric | int32 | `5` | INT(32, signed) | 4 byte little-endian | +| Exact Numeric | int64 | `6` | INT(64, signed) | 8 byte little-endian | +| Double | double | `7` | DOUBLE | IEEE little-endian | +| Exact Numeric | decimal4 | `8` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | +| Exact Numeric | decimal8 | `9` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | +| Exact Numeric | decimal16 | `10` | DECIMAL(precision, scale) | 1 byte scale in range [0, 38], followed by little-endian unscaled value (see decimal table) | +| Date | date | `11` | DATE | 4 byte little-endian | +| Timestamp | timestamp | `12` | TIMESTAMP(true, MICROS) | 8-byte little-endian | +| TimestampNTZ | timestamp without time zone | `13` | TIMESTAMP(false, MICROS) | 8-byte little-endian | +| Float | float | `14` | FLOAT | IEEE little-endian | +| Binary | binary | `15` | BINARY | 4 byte little-endian size, followed by bytes | +| String | string | `16` | STRING | 4 byte little-endian size, followed by UTF-8 encoded bytes | +| YMInterval | year-month interval | `19` | INT(32, signed)1 | 1 byte denoting start field (1 bit) and end field (1 bit) starting at LSB followed by 4-byte little-endian value. | +| DTInterval | day-time interval | `20` | INT(64, signed)1 | 1 byte denoting start field (2 bits) and end field (2 bits) starting at LSB followed by 8-byte little-endian value. | | Decimal Precision | Decimal value type | |-----------------------|--------------------| @@ -362,6 +362,8 @@ The Decimal type contains a scale, but no precision. The implied precision of a | 18 <= precision <= 38 | int128 | | > 38 | Not supported | +The *Logical Type* column indicates logical equivalence of physically encoded types. For example, a user expression operating on a string value containing "hello" should behave the same, whether it is encoded with the short string optimization, or long string encoding. Similarly, user expressions operating on an *int8* value of 1 should behave the same as a decimal16 with scale 2 and unscaled value 100. + The year-month and day-time interval types have one byte at the beginning indicating the start and end fields. In the case of the year-month interval, the least significant bit denotes the start field and the next least significant bit denotes the end field. The remaining 6 bits are unused. A field value of 0 represents YEAR and 1 represents MONTH. In the case of the day-time interval, the least significant 2 bits denote the start field and the next least significant 2 bits denote the end field. The remaining 4 bits are unused. A field value of 0 represents DAY, 1 represents HOUR, 2 represents MINUTE, and 3 represents SECOND. Type IDs 17 and 18 were originally reserved for a prototype feature (string-from-metadata) that was never implemented. These IDs are available for use by new types. From bc54eac57e5c69e31d2a1e6b0afefbfbf34b75a8 Mon Sep 17 00:00:00 2001 From: ivanjevtic-db Date: Thu, 12 Sep 2024 16:42:31 +0200 Subject: [PATCH 009/189] [SPARK-42846][SQL] Remove error condition _LEGACY_ERROR_TEMP_2011 ### What changes were proposed in this pull request? Removed error condition **_LEGACY_ERROR_TEMP_2011**, removed **dataTypeUnexpectedError**, **typeUnsupportedError**, and replaced them with **SparkException.internalError**. ### Why are the changes needed? It is impossible to trigger the error from user space. [Here](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-688ac8011f7fb514154ff57cfb1278b15aec481d68c1a499c90f8a330d3a42a1L141) I changed dataTypeUnexpectedError to internalError, since _typeSoFar_ argument will always be either: - NullType if this is the first row or - Some type which is returned by the [_inferField_](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-688ac8011f7fb514154ff57cfb1278b15aec481d68c1a499c90f8a330d3a42a1L125) function(which is a valid type). [Here](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-e9a88a888c1543c718c24f25036307cb32348ca3f618a8fa19240bdc3c0ffaf4L553) I changed typeUnsupportedError to internalError, since: - in [this](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-e9a88a888c1543c718c24f25036307cb32348ca3f618a8fa19240bdc3c0ffaf4L204) function call, the exception will be caught and - in [this](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-e9a88a888c1543c718c24f25036307cb32348ca3f618a8fa19240bdc3c0ffaf4L367) function call, a valid _desiredType_ will always be passed. [Here](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-d1fa4a2cbd66cff7d7d8a90d7ac70457a31e906cebb7d43a46a6036507fb4e7bL192) and [here](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-d1fa4a2cbd66cff7d7d8a90d7ac70457a31e906cebb7d43a46a6036507fb4e7bL217) I changed dataTypeUnexpectedError to internalError, since there is a type signature [here](https://github.com/apache/spark/compare/master...ivanjevtic-db:spark:remove-legacy-error-temp-2011?expand=1#diff-d1fa4a2cbd66cff7d7d8a90d7ac70457a31e906cebb7d43a46a6036507fb4e7bL104) which prevents _child.dataType_ from being unexpected type. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Build passing. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48086 from ivanjevtic-db/remove-legacy-error-temp-2011. Authored-by: ivanjevtic-db Signed-off-by: Max Gekk --- .../src/main/resources/error/error-conditions.json | 5 ----- .../spark/sql/catalyst/csv/CSVInferSchema.scala | 4 ++-- .../aggregate/ApproximatePercentile.scala | 6 +++--- .../spark/sql/errors/QueryExecutionErrors.scala | 12 ------------ .../execution/datasources/PartitioningUtils.scala | 4 ++-- 5 files changed, 7 insertions(+), 24 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 29eda228c2daa..0ebeea9aed8d2 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6656,11 +6656,6 @@ "Type does not support ordered operations." ] }, - "_LEGACY_ERROR_TEMP_2011" : { - "message" : [ - "Unexpected data type ." - ] - }, "_LEGACY_ERROR_TEMP_2013" : { "message" : [ "Negative values found in " diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVInferSchema.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVInferSchema.scala index 2c27da3cf6e15..5444ab6845867 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVInferSchema.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVInferSchema.scala @@ -21,12 +21,12 @@ import java.util.Locale import scala.util.control.Exception.allCatch +import org.apache.spark.SparkException import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.analysis.TypeCoercion import org.apache.spark.sql.catalyst.expressions.ExprUtils import org.apache.spark.sql.catalyst.util.{DateFormatter, TimestampFormatter} import org.apache.spark.sql.catalyst.util.LegacyDateFormats.FAST_DATE_FORMAT -import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -138,7 +138,7 @@ class CSVInferSchema(val options: CSVOptions) extends Serializable { case BooleanType => tryParseBoolean(field) case StringType => StringType case other: DataType => - throw QueryExecutionErrors.dataTypeUnexpectedError(other) + throw SparkException.internalError(s"Unexpected data type $other") } compatibleType(typeSoFar, typeElemInfer).getOrElse(StringType) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala index 4987e31b49911..8ad062ab0e2f9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala @@ -21,6 +21,7 @@ import java.nio.ByteBuffer import com.google.common.primitives.{Doubles, Ints, Longs} +import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, TypeCheckResult} import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess} @@ -32,7 +33,6 @@ import org.apache.spark.sql.catalyst.types.PhysicalNumericType import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData} import org.apache.spark.sql.catalyst.util.QuantileSummaries import org.apache.spark.sql.catalyst.util.QuantileSummaries.{defaultCompressThreshold, Stats} -import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.types._ import org.apache.spark.util.ArrayImplicits._ @@ -189,7 +189,7 @@ case class ApproximatePercentile( PhysicalNumericType.numeric(n) .toDouble(value.asInstanceOf[PhysicalNumericType#InternalType]) case other: DataType => - throw QueryExecutionErrors.dataTypeUnexpectedError(other) + throw SparkException.internalError(s"Unexpected data type $other") } buffer.add(doubleValue) } @@ -214,7 +214,7 @@ case class ApproximatePercentile( case DoubleType => doubleResult case _: DecimalType => doubleResult.map(Decimal(_)) case other: DataType => - throw QueryExecutionErrors.dataTypeUnexpectedError(other) + throw SparkException.internalError(s"Unexpected data type $other") } if (result.length == 0) { null diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 0b37cf951a29b..2ab86a5c5f03f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -384,18 +384,6 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE s"The aggregate window function ${toSQLId(funcName)} does not support merging.") } - def dataTypeUnexpectedError(dataType: DataType): SparkUnsupportedOperationException = { - new SparkUnsupportedOperationException( - errorClass = "_LEGACY_ERROR_TEMP_2011", - messageParameters = Map("dataType" -> dataType.catalogString)) - } - - def typeUnsupportedError(dataType: DataType): SparkIllegalArgumentException = { - new SparkIllegalArgumentException( - errorClass = "_LEGACY_ERROR_TEMP_2011", - messageParameters = Map("dataType" -> dataType.toString())) - } - def negativeValueUnexpectedError( frequencyExpression : Expression): SparkIllegalArgumentException = { new SparkIllegalArgumentException( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala index 676a2ab64d0a3..ffdca65151052 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala @@ -29,7 +29,7 @@ import scala.util.control.NonFatal import org.apache.hadoop.fs.Path -import org.apache.spark.SparkRuntimeException +import org.apache.spark.{SparkException, SparkRuntimeException} import org.apache.spark.sql.catalyst.{InternalRow, SQLConfHelper} import org.apache.spark.sql.catalyst.analysis.TypeCoercion import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec @@ -550,7 +550,7 @@ object PartitioningUtils extends SQLConfHelper { Cast(Literal(unescapePathName(value)), it).eval() case BinaryType => value.getBytes() case BooleanType => value.toBoolean - case dt => throw QueryExecutionErrors.typeUnsupportedError(dt) + case dt => throw SparkException.internalError(s"Unsupported partition type: $dt") } def validatePartitionColumn( From 1a0791d006e25898b67cc17e1420f053a39091b9 Mon Sep 17 00:00:00 2001 From: Bruce Robbins Date: Thu, 12 Sep 2024 08:11:03 -0700 Subject: [PATCH 010/189] [SPARK-49261][SQL] Don't replace literals in aggregate expressions with group-by expressions ### What changes were proposed in this pull request? Before this PR, `RewriteDistinctAggregates` could potentially replace literals in the aggregate expressions with output attributes from the `Expand` operator. This can occur when a group-by expression is a literal that happens by chance to match a literal used in an aggregate expression. E.g.: ``` create or replace temp view v1(a, b, c) as values (1, 1.001d, 2), (2, 3.001d, 4), (2, 3.001, 4); cache table v1; select round(sum(b), 6) as sum1, count(distinct a) as count1, count(distinct c) as count2 from ( select 6 as gb, * from v1 ) group by a, gb; ``` In the optimized plan, you can see that the literal 6 in the `round` function invocation has been patched with an output attribute (6#163) from the `Expand` operator: ``` == Optimized Logical Plan == 'Aggregate [a#123, 6#163], [round(first(sum(__auto_generated_subquery_name.b)#167, true) FILTER (WHERE (gid#162 = 0)), 6#163) AS sum1#114, count(__auto_generated_subquery_name.a#164) FILTER (WHERE (gid#162 = 1)) AS count1#115L, count(__auto_generated_subquery_name.c#165) FILTER (WHERE (gid#162 = 2)) AS count2#116L] +- Aggregate [a#123, 6#163, __auto_generated_subquery_name.a#164, __auto_generated_subquery_name.c#165, gid#162], [a#123, 6#163, __auto_generated_subquery_name.a#164, __auto_generated_subquery_name.c#165, gid#162, sum(__auto_generated_subquery_name.b#166) AS sum(__auto_generated_subquery_name.b)#167] +- Expand [[a#123, 6, null, null, 0, b#124], [a#123, 6, a#123, null, 1, null], [a#123, 6, null, c#125, 2, null]], [a#123, 6#163, __auto_generated_subquery_name.a#164, __auto_generated_subquery_name.c#165, gid#162, __auto_generated_subquery_name.b#166] +- InMemoryRelation [a#123, b#124, c#125], StorageLevel(disk, memory, deserialized, 1 replicas) +- LocalTableScan [a#6, b#7, c#8] ``` This is because the literal 6 was used in the group-by expressions (referred to as gb in the query, and renamed 6#163 in the `Expand` operator's output attributes). After this PR, foldable expressions in the aggregate expressions are kept as-is. ### Why are the changes needed? Some expressions require a foldable argument. In the above example, the `round` function requires a foldable expression as the scale argument. Because the scale argument is patched with an attribute, `RoundBase#checkInputDataTypes` returns an error, which leaves the `Aggregate` operator unresolved: ``` [INTERNAL_ERROR] Invalid call to dataType on unresolved object SQLSTATE: XX000 org.apache.spark.sql.catalyst.analysis.UnresolvedException: [INTERNAL_ERROR] Invalid call to dataType on unresolved object SQLSTATE: XX000 at org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute.dataType(unresolved.scala:255) at org.apache.spark.sql.catalyst.types.DataTypeUtils$.$anonfun$fromAttributes$1(DataTypeUtils.scala:241) at scala.collection.immutable.List.map(List.scala:247) at scala.collection.immutable.List.map(List.scala:79) at org.apache.spark.sql.catalyst.types.DataTypeUtils$.fromAttributes(DataTypeUtils.scala:241) at org.apache.spark.sql.catalyst.plans.QueryPlan.schema$lzycompute(QueryPlan.scala:428) at org.apache.spark.sql.catalyst.plans.QueryPlan.schema(QueryPlan.scala:428) at org.apache.spark.sql.execution.SparkPlan.executeCollectPublic(SparkPlan.scala:474) ... ``` ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? New tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47876 from bersprockets/group_by_lit_issue. Authored-by: Bruce Robbins Signed-off-by: Dongjoon Hyun --- .../optimizer/RewriteDistinctAggregates.scala | 3 ++- .../RewriteDistinctAggregatesSuite.scala | 18 +++++++++++++++- .../spark/sql/DataFrameAggregateSuite.scala | 21 +++++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala index 801bd2693af42..5aef82b64ed32 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala @@ -400,13 +400,14 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] { (distinctAggOperatorMap.flatMap(_._2) ++ regularAggOperatorMap.map(e => (e._1, e._3))).toMap + val groupByMapNonFoldable = groupByMap.filter(!_._1.foldable) val patchedAggExpressions = a.aggregateExpressions.map { e => e.transformDown { case e: Expression => // The same GROUP BY clauses can have different forms (different names for instance) in // the groupBy and aggregate expressions of an aggregate. This makes a map lookup // tricky. So we do a linear search for a semantically equal group by expression. - groupByMap + groupByMapNonFoldable .find(ge => e.semanticEquals(ge._1)) .map(_._2) .getOrElse(transformations.getOrElse(e, e)) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala index ac136dfb898ef..4d31999ded655 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.catalyst.expressions.{Literal, Round} import org.apache.spark.sql.catalyst.expressions.aggregate.CollectSet import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LocalRelation, LogicalPlan} @@ -109,4 +109,20 @@ class RewriteDistinctAggregatesSuite extends PlanTest { case _ => fail(s"Plan is not rewritten:\n$rewrite") } } + + test("SPARK-49261: Literals in grouping expressions shouldn't result in unresolved aggregation") { + val relation = testRelation2 + .select(Literal(6).as("gb"), $"a", $"b", $"c", $"d") + val input = relation + .groupBy($"a", $"gb")( + countDistinct($"b").as("agg1"), + countDistinct($"d").as("agg2"), + Round(sum($"c").as("sum1"), 6)).analyze + val rewriteFold = FoldablePropagation(input) + // without the fix, the below produces an unresolved plan + val rewrite = RewriteDistinctAggregates(rewriteFold) + if (!rewrite.resolved) { + fail(s"Plan is not as expected:\n$rewrite") + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 0e9d34c3bd96a..e80c3b23a7db3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -2490,6 +2490,27 @@ class DataFrameAggregateSuite extends QueryTest }) } } + + test("SPARK-49261: Literals in grouping expressions shouldn't result in unresolved aggregation") { + val data = Seq((1, 1.001d, 2), (2, 3.001d, 4), (2, 3.001, 4)).toDF("a", "b", "c") + withTempView("v1") { + data.createOrReplaceTempView("v1") + val df = + sql("""SELECT + | ROUND(SUM(b), 6) AS sum1, + | COUNT(DISTINCT a) AS count1, + | COUNT(DISTINCT c) AS count2 + |FROM ( + | SELECT + | 6 AS gb, + | * + | FROM v1 + |) + |GROUP BY a, gb + |""".stripMargin) + checkAnswer(df, Row(1.001d, 1, 1) :: Row(6.002d, 1, 1) :: Nil) + } + } } case class B(c: Option[Double]) From 1f24b2d72ed6821a6cc6d1d22683d2f3ba2326a2 Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Thu, 12 Sep 2024 09:26:56 -0700 Subject: [PATCH 011/189] [SPARK-44811][BUILD] Upgrade Guava to 33.2.1-jre ### What changes were proposed in this pull request? This PR upgrades Spark's built-in Guava from 14 to 33.2.1-jre Currently, Spark uses Guava 14 because the previous built-in Hive 2.3.9 is incompatible with new Guava versions. HIVE-27560 (https://github.com/apache/hive/pull/4542) makes Hive 2.3.10 compatible with Guava 14+ (thanks to LuciferYang) ### Why are the changes needed? It's a long-standing issue, see prior discussions at https://github.com/apache/spark/pull/35584, https://github.com/apache/spark/pull/36231, and https://github.com/apache/spark/pull/33989 ### Does this PR introduce _any_ user-facing change? Yes, some user-faced error messages changed. ### How was this patch tested? GA passed. Closes #42493 from pan3793/guava. Authored-by: Cheng Pan Signed-off-by: Dongjoon Hyun --- assembly/pom.xml | 2 +- core/pom.xml | 1 + dev/deps/spark-deps-hadoop-3-hive-2.3 | 7 ++++++- pom.xml | 3 ++- project/SparkBuild.scala | 2 +- .../catalyst/expressions/IntervalExpressionsSuite.scala | 2 +- .../test/resources/sql-tests/results/ansi/interval.sql.out | 4 ++-- .../src/test/resources/sql-tests/results/interval.sql.out | 4 ++-- 8 files changed, 16 insertions(+), 9 deletions(-) diff --git a/assembly/pom.xml b/assembly/pom.xml index 4b074a88dab4a..01bd324efc118 100644 --- a/assembly/pom.xml +++ b/assembly/pom.xml @@ -123,7 +123,7 @@ com.google.guava diff --git a/core/pom.xml b/core/pom.xml index 53d5ad71cebf5..19f58940ed942 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -558,6 +558,7 @@ org.eclipse.jetty:jetty-util org.eclipse.jetty:jetty-server com.google.guava:guava + com.google.guava:failureaccess com.google.protobuf:* diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index c89c92815d454..2db86ed229a01 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -33,6 +33,7 @@ breeze-macros_2.13/2.1.0//breeze-macros_2.13-2.1.0.jar breeze_2.13/2.1.0//breeze_2.13-2.1.0.jar bundle/2.24.6//bundle-2.24.6.jar cats-kernel_2.13/2.8.0//cats-kernel_2.13-2.8.0.jar +checker-qual/3.42.0//checker-qual-3.42.0.jar chill-java/0.10.0//chill-java-0.10.0.jar chill_2.13/0.10.0//chill_2.13-0.10.0.jar commons-cli/1.9.0//commons-cli-1.9.0.jar @@ -62,12 +63,14 @@ derby/10.16.1.1//derby-10.16.1.1.jar derbyshared/10.16.1.1//derbyshared-10.16.1.1.jar derbytools/10.16.1.1//derbytools-10.16.1.1.jar dropwizard-metrics-hadoop-metrics2-reporter/0.1.2//dropwizard-metrics-hadoop-metrics2-reporter-0.1.2.jar +error_prone_annotations/2.26.1//error_prone_annotations-2.26.1.jar esdk-obs-java/3.20.4.2//esdk-obs-java-3.20.4.2.jar +failureaccess/1.0.2//failureaccess-1.0.2.jar flatbuffers-java/24.3.25//flatbuffers-java-24.3.25.jar gcs-connector/hadoop3-2.2.21/shaded/gcs-connector-hadoop3-2.2.21-shaded.jar gmetric4j/1.0.10//gmetric4j-1.0.10.jar gson/2.11.0//gson-2.11.0.jar -guava/14.0.1//guava-14.0.1.jar +guava/33.2.1-jre//guava-33.2.1-jre.jar hadoop-aliyun/3.4.0//hadoop-aliyun-3.4.0.jar hadoop-annotations/3.4.0//hadoop-annotations-3.4.0.jar hadoop-aws/3.4.0//hadoop-aws-3.4.0.jar @@ -101,6 +104,7 @@ icu4j/75.1//icu4j-75.1.jar ini4j/0.5.4//ini4j-0.5.4.jar istack-commons-runtime/3.0.8//istack-commons-runtime-3.0.8.jar ivy/2.5.2//ivy-2.5.2.jar +j2objc-annotations/3.0.0//j2objc-annotations-3.0.0.jar jackson-annotations/2.17.2//jackson-annotations-2.17.2.jar jackson-core-asl/1.9.13//jackson-core-asl-1.9.13.jar jackson-core/2.17.2//jackson-core-2.17.2.jar @@ -184,6 +188,7 @@ lapack/3.0.3//lapack-3.0.3.jar leveldbjni-all/1.8//leveldbjni-all-1.8.jar libfb303/0.9.3//libfb303-0.9.3.jar libthrift/0.16.0//libthrift-0.16.0.jar +listenablefuture/9999.0-empty-to-avoid-conflict-with-guava//listenablefuture-9999.0-empty-to-avoid-conflict-with-guava.jar log4j-1.2-api/2.22.1//log4j-1.2-api-2.22.1.jar log4j-api/2.22.1//log4j-api-2.22.1.jar log4j-core/2.22.1//log4j-core-2.22.1.jar diff --git a/pom.xml b/pom.xml index 6f5c9b63f86de..b1497c7826855 100644 --- a/pom.xml +++ b/pom.xml @@ -195,7 +195,7 @@ 2.12.0 4.1.17 - 14.0.1 + 33.2.1-jre 2.11.0 3.1.9 3.0.12 @@ -3420,6 +3420,7 @@ org.spark-project.spark:unused com.google.guava:guava + com.google.guava:failureaccess org.jpmml:* diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index 01d4ad50a22ba..4a8214b2e20a3 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -1051,7 +1051,7 @@ object KubernetesIntegrationTests { * Overrides to work around sbt's dependency resolution being different from Maven's. */ object DependencyOverrides { - lazy val guavaVersion = sys.props.get("guava.version").getOrElse("14.0.1") + lazy val guavaVersion = sys.props.get("guava.version").getOrElse("33.1.0-jre") lazy val settings = Seq( dependencyOverrides += "com.google.guava" % "guava" % guavaVersion, dependencyOverrides += "xerces" % "xercesImpl" % "2.12.2", diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala index ff5ffe4e869a0..7caf23490a0ce 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/IntervalExpressionsSuite.scala @@ -351,7 +351,7 @@ class IntervalExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { Seq( (Period.ofMonths(2), Int.MaxValue) -> "overflow", - (Period.ofMonths(Int.MinValue), 10d) -> "not in range", + (Period.ofMonths(Int.MinValue), 10d) -> "out of range", (Period.ofMonths(-100), Float.NaN) -> "input is infinite or NaN", (Period.ofMonths(200), Double.PositiveInfinity) -> "input is infinite or NaN", (Period.ofMonths(-200), Float.NegativeInfinity) -> "input is infinite or NaN" diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/interval.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/interval.sql.out index 9e5c89045e514..b2f85835eb0df 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/interval.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/interval.sql.out @@ -2890,7 +2890,7 @@ SELECT (INTERVAL '-178956970-8' YEAR TO MONTH) / -1.0D struct<> -- !query output java.lang.ArithmeticException -not in range +rounded value is out of range for input 2.147483648E9 and rounding mode HALF_UP -- !query @@ -2970,7 +2970,7 @@ SELECT (INTERVAL '-106751991 04:00:54.775808' DAY TO SECOND) / -1.0D struct<> -- !query output java.lang.ArithmeticException -not in range +rounded value is out of range for input 9.223372036854776E18 and rounding mode HALF_UP -- !query diff --git a/sql/core/src/test/resources/sql-tests/results/interval.sql.out b/sql/core/src/test/resources/sql-tests/results/interval.sql.out index a4e1670517678..5471dafaec8eb 100644 --- a/sql/core/src/test/resources/sql-tests/results/interval.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/interval.sql.out @@ -2713,7 +2713,7 @@ SELECT (INTERVAL '-178956970-8' YEAR TO MONTH) / -1.0D struct<> -- !query output java.lang.ArithmeticException -not in range +rounded value is out of range for input 2.147483648E9 and rounding mode HALF_UP -- !query @@ -2793,7 +2793,7 @@ SELECT (INTERVAL '-106751991 04:00:54.775808' DAY TO SECOND) / -1.0D struct<> -- !query output java.lang.ArithmeticException -not in range +rounded value is out of range for input 9.223372036854776E18 and rounding mode HALF_UP -- !query From 98f0d9f32322074b01285f405c86df29997634a3 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Thu, 12 Sep 2024 18:52:43 +0200 Subject: [PATCH 012/189] [SPARK-49605][SQL] Fix the prompt when `ascendingOrder` is `DataTypeMismatch` in `SortArray` ### What changes were proposed in this pull request? The pr aims to fix the `prompt` when `ascendingOrder` is `DataTypeMismatch` in `SortArray`. ### Why are the changes needed? - Give an example with the following code: ```scala val df = Seq((Array[Int](2, 1, 3), true), (Array.empty[Int], false)).toDF("a", "b") df.selectExpr("sort_array(a, b)").collect() ``` - Before: ```scala scala> val df = Seq((Array[Int](2, 1, 3), true), (Array.empty[Int], false)).toDF("a", "b") val df: org.apache.spark.sql.DataFrame = [a: array, b: boolean] scala> df.selectExpr("sort_array(a, b)").collect() org.apache.spark.sql.catalyst.ExtendedAnalysisException: [DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "sort_array(a, b)" due to data type mismatch: The second parameter requires the "BOOLEAN" type, however "b" has the type "BOOLEAN". SQLSTATE: 42K09; line 1 pos 0; 'Project [unresolvedalias(sort_array(a#7, b#8))] +- Project [_1#2 AS a#7, _2#3 AS b#8] +- LocalRelation [_1#2, _2#3] at org.apache.spark.sql.catalyst.analysis.package$AnalysisErrorAt.dataTypeMismatch(package.scala:73) at org.apache.spark.sql.catalyst.analysis.CheckAnalysis.$anonfun$checkAnalysis0$7(CheckAnalysis.scala:331) at org.apache.spark.sql.catalyst.analysis.CheckAnalysis.$anonfun$checkAnalysis0$7$adapted(CheckAnalysis.scala:313) ``` image Obviously, this error message is `incorrect` and `confusing`. Through the following code: https://github.com/apache/spark/blob/8023504e69fdd037dea002e961b960fd9fa662ba/sql/api/src/main/scala/org/apache/spark/sql/functions.scala#L7176-L7195 we found that it actually requires `ascendingOrder` to be `foldable` and the data type to be `BooleanType`. - After: ``` scala> val df = Seq((Array[Int](2, 1, 3), true), (Array.empty[Int], false)).toDF("a", "b") val df: org.apache.spark.sql.DataFrame = [a: array, b: boolean] scala> df.selectExpr("sort_array(a, b)").collect() org.apache.spark.sql.catalyst.ExtendedAnalysisException: [DATATYPE_MISMATCH.NON_FOLDABLE_INPUT] Cannot resolve "sort_array(a, b)" due to data type mismatch: the input `ascendingOrder` should be a foldable "BOOLEAN" expression; however, got "b". SQLSTATE: 42K09; line 1 pos 0; 'Project [unresolvedalias(sort_array(a#7, b#8))] +- Project [_1#2 AS a#7, _2#3 AS b#8] +- LocalRelation [_1#2, _2#3] at org.apache.spark.sql.catalyst.analysis.package$AnalysisErrorAt.dataTypeMismatch(package.scala:73) at org.apache.spark.sql.catalyst.analysis.CheckAnalysis.$anonfun$checkAnalysis0$7(CheckAnalysis.scala:331) at org.apache.spark.sql.catalyst.analysis.CheckAnalysis.$anonfun$checkAnalysis0$7$adapted(CheckAnalysis.scala:313) ``` image ### Does this PR introduce _any_ user-facing change? Yes, When the value `ascendingOrder` in `SortArray` is `DataTypeMismatch`, the prompt is more `accurate`. ### How was this patch tested? - Add new UT - Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48082 from panbingkun/SPARK-49605. Authored-by: panbingkun Signed-off-by: Max Gekk --- .../expressions/collectionOperations.scala | 32 +++++++++++-------- .../analyzer-results/ansi/array.sql.out | 21 ++---------- .../sql-tests/analyzer-results/array.sql.out | 21 ++---------- .../sql-tests/results/ansi/array.sql.out | 22 ++----------- .../resources/sql-tests/results/array.sql.out | 22 ++----------- .../spark/sql/DataFrameFunctionsSuite.scala | 30 +++++++++++++++++ 6 files changed, 57 insertions(+), 91 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala index 5d5aece35383e..5cdd3c7eb62d1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala @@ -1058,20 +1058,26 @@ case class SortArray(base: Expression, ascendingOrder: Expression) override def checkInputDataTypes(): TypeCheckResult = base.dataType match { case ArrayType(dt, _) if RowOrdering.isOrderable(dt) => - ascendingOrder match { - case Literal(_: Boolean, BooleanType) => - TypeCheckResult.TypeCheckSuccess - case _ => - DataTypeMismatch( - errorSubClass = "UNEXPECTED_INPUT_TYPE", - messageParameters = Map( - "paramIndex" -> ordinalNumber(1), - "requiredType" -> toSQLType(BooleanType), - "inputSql" -> toSQLExpr(ascendingOrder), - "inputType" -> toSQLType(ascendingOrder.dataType)) - ) + if (!ascendingOrder.foldable) { + DataTypeMismatch( + errorSubClass = "NON_FOLDABLE_INPUT", + messageParameters = Map( + "inputName" -> toSQLId("ascendingOrder"), + "inputType" -> toSQLType(ascendingOrder.dataType), + "inputExpr" -> toSQLExpr(ascendingOrder))) + } else if (ascendingOrder.dataType != BooleanType) { + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(1), + "requiredType" -> toSQLType(BooleanType), + "inputSql" -> toSQLExpr(ascendingOrder), + "inputType" -> toSQLType(ascendingOrder.dataType)) + ) + } else { + TypeCheckResult.TypeCheckSuccess } - case ArrayType(dt, _) => + case ArrayType(_, _) => DataTypeMismatch( errorSubClass = "INVALID_ORDERING_TYPE", messageParameters = Map( diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/array.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/array.sql.out index 57108c4582f45..53595d1b8a3eb 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/array.sql.out @@ -194,25 +194,8 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query select sort_array(array('b', 'd'), cast(NULL as boolean)) -- !query analysis -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", - "sqlState" : "42K09", - "messageParameters" : { - "inputSql" : "\"CAST(NULL AS BOOLEAN)\"", - "inputType" : "\"BOOLEAN\"", - "paramIndex" : "second", - "requiredType" : "\"BOOLEAN\"", - "sqlExpr" : "\"sort_array(array(b, d), CAST(NULL AS BOOLEAN))\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 57, - "fragment" : "sort_array(array('b', 'd'), cast(NULL as boolean))" - } ] -} +Project [sort_array(array(b, d), cast(null as boolean)) AS sort_array(array(b, d), CAST(NULL AS BOOLEAN))#x] ++- OneRowRelation -- !query diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out index fb331089d7545..4db56d6c70561 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out @@ -194,25 +194,8 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query select sort_array(array('b', 'd'), cast(NULL as boolean)) -- !query analysis -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", - "sqlState" : "42K09", - "messageParameters" : { - "inputSql" : "\"CAST(NULL AS BOOLEAN)\"", - "inputType" : "\"BOOLEAN\"", - "paramIndex" : "second", - "requiredType" : "\"BOOLEAN\"", - "sqlExpr" : "\"sort_array(array(b, d), CAST(NULL AS BOOLEAN))\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 57, - "fragment" : "sort_array(array('b', 'd'), cast(NULL as boolean))" - } ] -} +Project [sort_array(array(b, d), cast(null as boolean)) AS sort_array(array(b, d), CAST(NULL AS BOOLEAN))#x] ++- OneRowRelation -- !query diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/array.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/array.sql.out index d17d87900fc71..7394e428091c7 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/array.sql.out @@ -151,27 +151,9 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query select sort_array(array('b', 'd'), cast(NULL as boolean)) -- !query schema -struct<> +struct> -- !query output -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", - "sqlState" : "42K09", - "messageParameters" : { - "inputSql" : "\"CAST(NULL AS BOOLEAN)\"", - "inputType" : "\"BOOLEAN\"", - "paramIndex" : "second", - "requiredType" : "\"BOOLEAN\"", - "sqlExpr" : "\"sort_array(array(b, d), CAST(NULL AS BOOLEAN))\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 57, - "fragment" : "sort_array(array('b', 'd'), cast(NULL as boolean))" - } ] -} +NULL -- !query diff --git a/sql/core/src/test/resources/sql-tests/results/array.sql.out b/sql/core/src/test/resources/sql-tests/results/array.sql.out index 92da0a490ff81..c1330c620acfb 100644 --- a/sql/core/src/test/resources/sql-tests/results/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/array.sql.out @@ -151,27 +151,9 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query select sort_array(array('b', 'd'), cast(NULL as boolean)) -- !query schema -struct<> +struct> -- !query output -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", - "sqlState" : "42K09", - "messageParameters" : { - "inputSql" : "\"CAST(NULL AS BOOLEAN)\"", - "inputType" : "\"BOOLEAN\"", - "paramIndex" : "second", - "requiredType" : "\"BOOLEAN\"", - "sqlExpr" : "\"sort_array(array(b, d), CAST(NULL AS BOOLEAN))\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 57, - "fragment" : "sort_array(array('b', 'd'), cast(NULL as boolean))" - } ] -} +NULL -- !query diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index d488adc5ac3d1..f16171940df21 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala @@ -964,6 +964,36 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession { queryContext = Array(ExpectedContext("", "", 0, 12, "sort_array(a)")) ) + val df4 = Seq((Array[Int](2, 1, 3), true), (Array.empty[Int], false)).toDF("a", "b") + checkError( + exception = intercept[AnalysisException] { + df4.selectExpr("sort_array(a, b)").collect() + }, + condition = "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + sqlState = "42K09", + parameters = Map( + "inputName" -> "`ascendingOrder`", + "inputType" -> "\"BOOLEAN\"", + "inputExpr" -> "\"b\"", + "sqlExpr" -> "\"sort_array(a, b)\""), + context = ExpectedContext(fragment = "sort_array(a, b)", start = 0, stop = 15) + ) + + checkError( + exception = intercept[AnalysisException] { + df4.selectExpr("sort_array(a, 'A')").collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = "42K09", + parameters = Map( + "sqlExpr" -> "\"sort_array(a, A)\"", + "paramIndex" -> "second", + "inputSql" -> "\"A\"", + "inputType" -> "\"STRING\"", + "requiredType" -> "\"BOOLEAN\""), + context = ExpectedContext(fragment = "sort_array(a, 'A')", start = 0, stop = 17) + ) + checkAnswer( df.select(array_sort($"a"), array_sort($"b")), Seq( From 317eddb7390c9b3b836108b6ffa65110b6163c33 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Thu, 12 Sep 2024 10:57:31 -0700 Subject: [PATCH 013/189] [SPARK-49606][PS][DOCS] Improve documentation of Pandas on Spark plotting API ### What changes were proposed in this pull request? Improve documentation of Pandas on Spark plotting API following pandas 2.2 (stable), see https://pandas.pydata.org/docs/reference/frame.html. ### Why are the changes needed? Better documentation and parity with pandas. ### Does this PR introduce _any_ user-facing change? Doc changes only. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48083 from xinrong-meng/doc_impr. Authored-by: Xinrong Meng Signed-off-by: Dongjoon Hyun --- python/pyspark/pandas/plot/core.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index 453b17834020e..067c7db664dee 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -540,7 +540,7 @@ def line(self, x=None, y=None, **kwargs): """ Plot DataFrame/Series as lines. - This function is useful to plot lines using Series's values + This function is useful to plot lines using DataFrame’s values as coordinates. Parameters @@ -606,6 +606,12 @@ def bar(self, x=None, y=None, **kwds): """ Vertical bar plot. + A bar plot is a plot that presents categorical data with rectangular + bars with lengths proportional to the values that they represent. A + bar plot shows comparisons among discrete categories. One axis of the + plot shows the specific categories being compared, and the other axis + represents a measured value. + Parameters ---------- x : label or position, optional @@ -797,7 +803,17 @@ def barh(self, x=None, y=None, **kwargs): def box(self, **kwds): """ - Make a box plot of the Series columns. + Make a box plot of the DataFrame columns. + + A box plot is a method for graphically depicting groups of numerical data through + their quartiles. The box extends from the Q1 to Q3 quartile values of the data, + with a line at the median (Q2). The whiskers extend from the edges of box to show + the range of the data. The position of the whiskers is set by default to + 1.5*IQR (IQR = Q3 - Q1) from the edges of the box. Outlier points are those past + the end of the whiskers. + + A consideration when using this chart is that the box and the whiskers can overlap, + which is very common when plotting small sets of data. Parameters ---------- @@ -851,9 +867,11 @@ def box(self, **kwds): def hist(self, bins=10, **kwds): """ Draw one histogram of the DataFrame’s columns. + A `histogram`_ is a representation of the distribution of data. This function calls :meth:`plotting.backend.plot`, on each series in the DataFrame, resulting in one histogram per column. + This is useful when the DataFrame’s Series are in a similar scale. .. _histogram: https://en.wikipedia.org/wiki/Histogram @@ -902,6 +920,10 @@ def kde(self, bw_method=None, ind=None, **kwargs): """ Generate Kernel Density Estimate plot using Gaussian kernels. + In statistics, kernel density estimation (KDE) is a non-parametric way to + estimate the probability density function (PDF) of a random variable. This + function uses Gaussian kernels and includes automatic bandwidth determination. + Parameters ---------- bw_method : scalar From d2d293e3fb57d6c9dea084b5fe6707d67c715af3 Mon Sep 17 00:00:00 2001 From: prathit06 Date: Thu, 12 Sep 2024 12:25:17 -0700 Subject: [PATCH 014/189] [SPARK-49598][K8S] Support user-defined labels for OnDemand PVCs ### What changes were proposed in this pull request? Currently when user sets `volumes.persistentVolumeClaim.[VolumeName].options.claimName=OnDemand` PVCs are created with only 1 label i.e. spark-app-selector = spark.app.id. Objective of this PR is to allow support of custom labels for onDemand PVCs ### Why are the changes needed? Changes are needed so users can set custom labels to PVCs ### Does this PR introduce _any_ user-facing change? It does not break any existing behaviour but adds a new feature/improvement to enable custom label additions in ondemand PVCs ### How was this patch tested? This was tested in internal/production k8 cluster ### Was this patch authored or co-authored using generative AI tooling? No Closes #48079 from prathit06/ondemand-pvc-labels. Lead-authored-by: prathit06 Co-authored-by: Prathit malik <53890994+prathit06@users.noreply.github.com> Signed-off-by: Dongjoon Hyun --- docs/running-on-kubernetes.md | 18 +++++ .../org/apache/spark/deploy/k8s/Config.scala | 2 +- .../deploy/k8s/KubernetesVolumeSpec.scala | 3 +- .../deploy/k8s/KubernetesVolumeUtils.scala | 16 +++- .../features/MountVolumesFeatureStep.scala | 9 ++- .../spark/deploy/k8s/KubernetesTestConf.scala | 9 ++- .../k8s/KubernetesVolumeUtilsSuite.scala | 34 ++++++++- .../MountVolumesFeatureStepSuite.scala | 73 +++++++++++++++++++ 8 files changed, 154 insertions(+), 10 deletions(-) diff --git a/docs/running-on-kubernetes.md b/docs/running-on-kubernetes.md index c3c567e1b8224..d8be32e047717 100644 --- a/docs/running-on-kubernetes.md +++ b/docs/running-on-kubernetes.md @@ -1182,6 +1182,15 @@ See the [configuration page](configuration.html) for information on Spark config 2.4.0 + + spark.kubernetes.driver.volumes.[VolumeType].[VolumeName].label.[LabelName] + (none) + + Configure Kubernetes Volume labels passed to the Kubernetes with LabelName as key having specified value, must conform with Kubernetes label format. For example, + spark.kubernetes.driver.volumes.persistentVolumeClaim.checkpointpvc.label.foo=bar. + + 4.0.0 + spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].mount.path (none) @@ -1218,6 +1227,15 @@ See the [configuration page](configuration.html) for information on Spark config 2.4.0 + + spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].label.[LabelName] + (none) + + Configure Kubernetes Volume labels passed to the Kubernetes with LabelName as key having specified value, must conform with Kubernetes label format. For example, + spark.kubernetes.executor.volumes.persistentVolumeClaim.checkpointpvc.label.foo=bar. + + 4.0.0 + spark.kubernetes.local.dirs.tmpfs false diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala index 393ffc5674011..3a4d68c19014d 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala @@ -776,7 +776,7 @@ private[spark] object Config extends Logging { val KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY = "options.medium" val KUBERNETES_VOLUMES_OPTIONS_SIZE_LIMIT_KEY = "options.sizeLimit" val KUBERNETES_VOLUMES_OPTIONS_SERVER_KEY = "options.server" - + val KUBERNETES_VOLUMES_LABEL_KEY = "label." val KUBERNETES_DRIVER_ENV_PREFIX = "spark.kubernetes.driverEnv." val KUBERNETES_DNS_SUBDOMAIN_NAME_MAX_LENGTH = 253 diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala index 3f7355de18911..9dfd40a773eb1 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala @@ -24,7 +24,8 @@ private[spark] case class KubernetesHostPathVolumeConf(hostPath: String) private[spark] case class KubernetesPVCVolumeConf( claimName: String, storageClass: Option[String] = None, - size: Option[String] = None) + size: Option[String] = None, + labels: Option[Map[String, String]] = None) extends KubernetesVolumeSpecificConf private[spark] case class KubernetesEmptyDirVolumeConf( diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala index ee2108e8234d3..6463512c0114b 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala @@ -45,13 +45,21 @@ object KubernetesVolumeUtils { val pathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_PATH_KEY" val readOnlyKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_READONLY_KEY" val subPathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_SUBPATH_KEY" + val labelKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_LABEL_KEY" + + val volumeLabelsMap = properties + .filter(_._1.startsWith(labelKey)) + .map { + case (k, v) => k.replaceAll(labelKey, "") -> v + } KubernetesVolumeSpec( volumeName = volumeName, mountPath = properties(pathKey), mountSubPath = properties.getOrElse(subPathKey, ""), mountReadOnly = properties.get(readOnlyKey).exists(_.toBoolean), - volumeConf = parseVolumeSpecificConf(properties, volumeType, volumeName)) + volumeConf = parseVolumeSpecificConf(properties, + volumeType, volumeName, Option(volumeLabelsMap))) }.toSeq } @@ -74,7 +82,8 @@ object KubernetesVolumeUtils { private def parseVolumeSpecificConf( options: Map[String, String], volumeType: String, - volumeName: String): KubernetesVolumeSpecificConf = { + volumeName: String, + labels: Option[Map[String, String]]): KubernetesVolumeSpecificConf = { volumeType match { case KUBERNETES_VOLUMES_HOSTPATH_TYPE => val pathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_PATH_KEY" @@ -91,7 +100,8 @@ object KubernetesVolumeUtils { KubernetesPVCVolumeConf( options(claimNameKey), options.get(storageClassKey), - options.get(sizeLimitKey)) + options.get(sizeLimitKey), + labels) case KUBERNETES_VOLUMES_EMPTYDIR_TYPE => val mediumKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY" diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala index 72cc012a6bdd0..5cc61c746b0e0 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala @@ -74,7 +74,7 @@ private[spark] class MountVolumesFeatureStep(conf: KubernetesConf) new VolumeBuilder() .withHostPath(new HostPathVolumeSource(hostPath, "")) - case KubernetesPVCVolumeConf(claimNameTemplate, storageClass, size) => + case KubernetesPVCVolumeConf(claimNameTemplate, storageClass, size, labels) => val claimName = conf match { case c: KubernetesExecutorConf => claimNameTemplate @@ -86,12 +86,17 @@ private[spark] class MountVolumesFeatureStep(conf: KubernetesConf) .replaceAll(PVC_ON_DEMAND, s"${conf.resourceNamePrefix}-driver$PVC_POSTFIX-$i") } if (storageClass.isDefined && size.isDefined) { + val defaultVolumeLabels = Map(SPARK_APP_ID_LABEL -> conf.appId) + val volumeLabels = labels match { + case Some(customLabelsMap) => (customLabelsMap ++ defaultVolumeLabels).asJava + case None => defaultVolumeLabels.asJava + } additionalResources.append(new PersistentVolumeClaimBuilder() .withKind(PVC) .withApiVersion("v1") .withNewMetadata() .withName(claimName) - .addToLabels(SPARK_APP_ID_LABEL, conf.appId) + .addToLabels(volumeLabels) .endMetadata() .withNewSpec() .withStorageClassName(storageClass.get) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala index b70b9348d23b4..7e0a65bcdda90 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala @@ -117,12 +117,17 @@ object KubernetesTestConf { (KUBERNETES_VOLUMES_HOSTPATH_TYPE, Map(KUBERNETES_VOLUMES_OPTIONS_PATH_KEY -> path)) - case KubernetesPVCVolumeConf(claimName, storageClass, sizeLimit) => + case KubernetesPVCVolumeConf(claimName, storageClass, sizeLimit, labels) => val sconf = storageClass .map { s => (KUBERNETES_VOLUMES_OPTIONS_CLAIM_STORAGE_CLASS_KEY, s) }.toMap val lconf = sizeLimit.map { l => (KUBERNETES_VOLUMES_OPTIONS_SIZE_LIMIT_KEY, l) }.toMap + val llabels = labels match { + case Some(value) => value.map { case(k, v) => s"label.$k" -> v } + case None => Map() + } (KUBERNETES_VOLUMES_PVC_TYPE, - Map(KUBERNETES_VOLUMES_OPTIONS_CLAIM_NAME_KEY -> claimName) ++ sconf ++ lconf) + Map(KUBERNETES_VOLUMES_OPTIONS_CLAIM_NAME_KEY -> claimName) ++ + sconf ++ lconf ++ llabels) case KubernetesEmptyDirVolumeConf(medium, sizeLimit) => val mconf = medium.map { m => (KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY, m) }.toMap diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala index fdc1aae0d4109..5c103739d3082 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala @@ -56,7 +56,39 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.mountPath === "/path") assert(volumeSpec.mountReadOnly) assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === - KubernetesPVCVolumeConf("claimName")) + KubernetesPVCVolumeConf("claimName", labels = Some(Map()))) + } + + test("SPARK-49598: Parses persistentVolumeClaim volumes correctly with labels") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.path", "/path") + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.readOnly", "true") + sparkConf.set("test.persistentVolumeClaim.volumeName.options.claimName", "claimName") + sparkConf.set("test.persistentVolumeClaim.volumeName.label.env", "test") + sparkConf.set("test.persistentVolumeClaim.volumeName.label.foo", "bar") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.mountReadOnly) + assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === + KubernetesPVCVolumeConf(claimName = "claimName", + labels = Some(Map("env" -> "test", "foo" -> "bar")))) + } + + test("SPARK-49598: Parses persistentVolumeClaim volumes & puts " + + "labels as empty Map if not provided") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.path", "/path") + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.readOnly", "true") + sparkConf.set("test.persistentVolumeClaim.volumeName.options.claimName", "claimName") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.mountReadOnly) + assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === + KubernetesPVCVolumeConf(claimName = "claimName", labels = Some(Map()))) } test("Parses emptyDir volumes correctly") { diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala index 54796def95e53..6a68898c5f61c 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala @@ -131,6 +131,79 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { assert(pvcClaim.getClaimName.endsWith("-driver-pvc-0")) } + test("SPARK-49598: Create and mounts persistentVolumeClaims in driver with labels") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + true, + KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, + storageClass = Some("gp3"), + size = Some("1Mi"), + labels = Some(Map("foo" -> "bar", "env" -> "test"))) + ) + + val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + assert(configuredPod.pod.getSpec.getVolumes.size() === 1) + val pvcClaim = configuredPod.pod.getSpec.getVolumes.get(0).getPersistentVolumeClaim + assert(pvcClaim.getClaimName.endsWith("-driver-pvc-0")) + } + + test("SPARK-49598: Create and mounts persistentVolumeClaims in executors with labels") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + true, + KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, + storageClass = Some("gp3"), + size = Some("1Mi"), + labels = Some(Map("foo1" -> "bar1", "env" -> "exec-test"))) + ) + + val executorConf = KubernetesTestConf.createExecutorConf(volumes = Seq(volumeConf)) + val executorStep = new MountVolumesFeatureStep(executorConf) + val executorPod = executorStep.configurePod(SparkPod.initialPod()) + + assert(executorPod.pod.getSpec.getVolumes.size() === 1) + val executorPVC = executorPod.pod.getSpec.getVolumes.get(0).getPersistentVolumeClaim + assert(executorPVC.getClaimName.endsWith("-exec-1-pvc-0")) + } + + test("SPARK-49598: Mount multiple volumes to executor with labels") { + val pvcVolumeConf1 = KubernetesVolumeSpec( + "checkpointVolume1", + "/checkpoints1", + "", + true, + KubernetesPVCVolumeConf(claimName = "pvcClaim1", + storageClass = Some("gp3"), + size = Some("1Mi"), + labels = Some(Map("foo1" -> "bar1", "env1" -> "exec-test-1"))) + ) + + val pvcVolumeConf2 = KubernetesVolumeSpec( + "checkpointVolume2", + "/checkpoints2", + "", + true, + KubernetesPVCVolumeConf(claimName = "pvcClaim2", + storageClass = Some("gp3"), + size = Some("1Mi"), + labels = Some(Map("foo2" -> "bar2", "env2" -> "exec-test-2"))) + ) + + val kubernetesConf = KubernetesTestConf.createExecutorConf( + volumes = Seq(pvcVolumeConf1, pvcVolumeConf2)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + + assert(configuredPod.pod.getSpec.getVolumes.size() === 2) + assert(configuredPod.container.getVolumeMounts.size() === 2) + } + test("Create and mount persistentVolumeClaims in executors") { val volumeConf = KubernetesVolumeSpec( "testVolume", From 875def39166549f9de54b141f5397cb3f74a918e Mon Sep 17 00:00:00 2001 From: Wei Guo Date: Thu, 12 Sep 2024 16:26:33 -0700 Subject: [PATCH 015/189] [SPARK-49081][SQL][DOCS] Add data source options docs of `Protobuf` ### What changes were proposed in this pull request? This PR aims to add data source options docs of `Protobuf` data source. Other data sources such as `csv`, `json` have corresponding options documents. The document section appears as follows: image image ### Why are the changes needed? In order to facilitate Spark users to better understand and use the options of `Protobuf` data source. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA and local manual check with `SKIP_API=1 bundle exec jekyll build --watch`. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47570 from wayneguow/pb_docs. Authored-by: Wei Guo Signed-off-by: Dongjoon Hyun --- .../sql/protobuf/utils/ProtobufOptions.scala | 12 ++-- docs/sql-data-sources-protobuf.md | 67 ++++++++++++++++++- 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala index 6644bce98293b..e85097a272f24 100644 --- a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala +++ b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufOptions.scala @@ -43,8 +43,8 @@ private[sql] class ProtobufOptions( /** * Adds support for recursive fields. If this option is is not specified, recursive fields are - * not permitted. Setting it to 0 drops the recursive fields, 1 allows it to be recursed once, - * and 2 allows it to be recursed twice and so on, up to 10. Values larger than 10 are not + * not permitted. Setting it to 1 drops the recursive fields, 0 allows it to be recursed once, + * and 3 allows it to be recursed twice and so on, up to 10. Values larger than 10 are not * allowed in order avoid inadvertently creating very large schemas. If a Protobuf message * has depth beyond this limit, the Spark struct returned is truncated after the recursion limit. * @@ -52,8 +52,8 @@ private[sql] class ProtobufOptions( * `message Person { string name = 1; Person friend = 2; }` * The following lists the schema with different values for this setting. * 1: `struct` - * 2: `struct>` - * 3: `struct>>` + * 2: `struct>` + * 3: `struct>>` * and so on. */ val recursiveFieldMaxDepth: Int = parameters.getOrElse("recursive.fields.max.depth", "-1").toInt @@ -181,7 +181,7 @@ private[sql] class ProtobufOptions( val upcastUnsignedInts: Boolean = parameters.getOrElse("upcast.unsigned.ints", false.toString).toBoolean - // Whether to unwrap the struct representation for well known primitve wrapper types when + // Whether to unwrap the struct representation for well known primitive wrapper types when // deserializing. By default, the wrapper types for primitives (i.e. google.protobuf.Int32Value, // google.protobuf.Int64Value, etc.) will get deserialized as structs. We allow the option to // deserialize them as their respective primitives. @@ -221,7 +221,7 @@ private[sql] class ProtobufOptions( // By default, in the spark schema field a will be dropped, which result in schema // b struct // If retain.empty.message.types=true, field a will be retained by inserting a dummy column. - // b struct, name: string> + // b struct, name: string> val retainEmptyMessage: Boolean = parameters.getOrElse("retain.empty.message.types", false.toString).toBoolean } diff --git a/docs/sql-data-sources-protobuf.md b/docs/sql-data-sources-protobuf.md index 34cb1d4997d28..4dd6579f92cd2 100644 --- a/docs/sql-data-sources-protobuf.md +++ b/docs/sql-data-sources-protobuf.md @@ -434,4 +434,69 @@ message Person {
```
- \ No newline at end of file + + +## Data Source Option + +Data source options of Protobuf can be set via: +* the built-in functions below + * `from_protobuf` + * `to_protobuf` + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Property NameDefaultMeaningScope
modeFAILFASTAllows a mode for dealing with corrupt records during parsing.
+
    +
  • PERMISSIVE: when it meets a corrupted record, sets all fields to null.
    • +
  • DROPMALFORMED: ignores the whole corrupted records. This mode is unsupported in the Protobuf built-in functions.
    • +
  • FAILFAST: throws an exception when it meets corrupted records.
    • +
+ 		read
recursive.fields.max.depth-1Specifies the maximum number of recursion levels to allow when parsing the schema. For more details refers to the section Handling circular references protobuf fields.read
convert.any.fields.to.jsonfalseEnables converting Protobuf Any fields to JSON. This option should be enabled carefully. JSON conversion and processing are inefficient. In addition, schema safety is also reduced making downstream processing error-prone.read
emit.default.valuesfalseWhether to render fields with zero values when deserializing Protobuf to a Spark struct. When a field is empty in the serialized Protobuf, this library will deserialize them as null by default, this option can control whether to render the type-specific zero values.read
enums.as.intsfalseWhether to render enum fields as their integer values. When this option set to false, an enum field will be mapped to StringType, and the value is the name of enum; when set to true, an enum field will be mapped to IntegerType, the value is its integer value.read
upcast.unsigned.intsfalseWhether to upcast unsigned integers into a larger type. Setting this option to true, LongType is used for uint32 and Decimal(20, 0) is used for uint64, so their representation can contain large unsigned values without overflow.read
unwrap.primitive.wrapper.typesfalseWhether to unwrap the struct representation for well-known primitive wrapper types when deserializing. By default, the wrapper types for primitives (i.e. google.protobuf.Int32Value, google.protobuf.Int64Value, etc.) will get deserialized as structs.read
retain.empty.message.typesfalseWhether to retain fields of the empty proto message type in Schema. Since Spark doesn't allow writing empty StructType, the empty proto message type will be dropped by default. Setting this option to true will insert a dummy column(__dummy_field_in_empty_struct) to the empty proto message so that the empty message fields will be retained.read
From 3b8dddac65bce6f88f51e23e777d521d65fa3373 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Fri, 13 Sep 2024 09:21:20 +0800 Subject: [PATCH 016/189] [SPARK-49531][PYTHON][CONNECT] Support line plot with plotly backend ### What changes were proposed in this pull request? Support line plot with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations, such as line plots, by leveraging libraries like Plotly. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. ```python >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] >>> columns = ["category", "int_val", "float_val"] >>> sdf = spark.createDataFrame(data, columns) >>> sdf.show() +--------+-------+---------+ |category|int_val|float_val| +--------+-------+---------+ | A| 10| 1.5| | B| 30| 2.5| | C| 20| 3.5| +--------+-------+---------+ >>> f = sdf.plot(kind="line", x="category", y="int_val") >>> f.show() # see below >>> g = sdf.plot.line(x="category", y=["int_val", "float_val"]) >>> g.show() # see below ``` `f.show()`: ![newplot](https://github.com/user-attachments/assets/ebd50bbc-0dd1-437f-ae0c-0b4de8f3c722) `g.show()`: ![newplot (1)](https://github.com/user-attachments/assets/46d28840-a147-428f-8d88-d424aa76ad06) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48008 from xinrong-meng/plot_line. Authored-by: Xinrong Meng Signed-off-by: Ruifeng Zheng --- dev/sparktestsupport/modules.py | 4 + python/pyspark/errors/error-conditions.json | 5 + python/pyspark/sql/classic/dataframe.py | 5 + python/pyspark/sql/connect/dataframe.py | 5 + python/pyspark/sql/dataframe.py | 27 ++++ python/pyspark/sql/plot/__init__.py | 21 +++ python/pyspark/sql/plot/core.py | 135 ++++++++++++++++++ python/pyspark/sql/plot/plotly.py | 30 ++++ .../tests/connect/test_parity_frame_plot.py | 36 +++++ .../connect/test_parity_frame_plot_plotly.py | 36 +++++ python/pyspark/sql/tests/plot/__init__.py | 16 +++ .../pyspark/sql/tests/plot/test_frame_plot.py | 79 ++++++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 64 +++++++++ python/pyspark/sql/utils.py | 17 +++ python/pyspark/testing/sqlutils.py | 7 + .../apache/spark/sql/internal/SQLConf.scala | 27 ++++ 16 files changed, 514 insertions(+) create mode 100644 python/pyspark/sql/plot/__init__.py create mode 100644 python/pyspark/sql/plot/core.py create mode 100644 python/pyspark/sql/plot/plotly.py create mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot.py create mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py create mode 100644 python/pyspark/sql/tests/plot/__init__.py create mode 100644 python/pyspark/sql/tests/plot/test_frame_plot.py create mode 100644 python/pyspark/sql/tests/plot/test_frame_plot_plotly.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 34fbb8450d544..b9a4bed715f67 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -548,6 +548,8 @@ def __hash__(self): "pyspark.sql.tests.test_udtf", "pyspark.sql.tests.test_utils", "pyspark.sql.tests.test_resources", + "pyspark.sql.tests.plot.test_frame_plot", + "pyspark.sql.tests.plot.test_frame_plot_plotly", ], ) @@ -1051,6 +1053,8 @@ def __hash__(self): "pyspark.sql.tests.connect.test_parity_arrow_cogrouped_map", "pyspark.sql.tests.connect.test_parity_python_datasource", "pyspark.sql.tests.connect.test_parity_python_streaming_datasource", + "pyspark.sql.tests.connect.test_parity_frame_plot", + "pyspark.sql.tests.connect.test_parity_frame_plot_plotly", "pyspark.sql.tests.connect.test_utils", "pyspark.sql.tests.connect.client.test_artifact", "pyspark.sql.tests.connect.client.test_artifact_localcluster", diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 4061d024a83cd..92aeb15e21d1b 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -1088,6 +1088,11 @@ "Function `` should use only POSITIONAL or POSITIONAL OR KEYWORD arguments." ] }, + "UNSUPPORTED_PLOT_BACKEND": { + "message": [ + "`` is not supported, it should be one of the values from " + ] + }, "UNSUPPORTED_SIGNATURE": { "message": [ "Unsupported signature: ." diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index 91b9591625904..d174f7774cc57 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -58,6 +58,7 @@ from pyspark.sql.classic.column import _to_seq, _to_list, _to_java_column from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter +from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.types import ( StructType, @@ -1862,6 +1863,10 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: messageParameters={"member": "queryExecution"}, ) + @property + def plot(self) -> PySparkPlotAccessor: + return PySparkPlotAccessor(self) + class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 768abd655d497..e3b1d35b2d5d6 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -83,6 +83,7 @@ UnresolvedStar, ) from pyspark.sql.connect.functions import builtin as F +from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema from pyspark.sql.pandas.functions import _validate_pandas_udf # type: ignore[attr-defined] @@ -2239,6 +2240,10 @@ def rdd(self) -> "RDD[Row]": def executionInfo(self) -> Optional["ExecutionInfo"]: return self._execution_info + @property + def plot(self) -> PySparkPlotAccessor: + return PySparkPlotAccessor(self) + class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index ef35b73332572..7748510258eaa 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -39,6 +39,7 @@ from pyspark.sql.column import Column from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter +from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.types import StructType, Row from pyspark.sql.utils import dispatch_df_method @@ -6394,6 +6395,32 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: """ ... + @property + def plot(self) -> PySparkPlotAccessor: + """ + Returns a :class:`PySparkPlotAccessor` for plotting functions. + + .. versionadded:: 4.0.0 + + Returns + ------- + :class:`PySparkPlotAccessor` + + Notes + ----- + This API is experimental. + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> type(df.plot) + + >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP + """ + ... + class DataFrameNaFunctions: """Functionality for working with missing data in :class:`DataFrame`. diff --git a/python/pyspark/sql/plot/__init__.py b/python/pyspark/sql/plot/__init__.py new file mode 100644 index 0000000000000..6da07061b2a09 --- /dev/null +++ b/python/pyspark/sql/plot/__init__.py @@ -0,0 +1,21 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +This package includes the plotting APIs for PySpark DataFrame. +""" +from pyspark.sql.plot.core import * # noqa: F403, F401 diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py new file mode 100644 index 0000000000000..baee610dc6bd0 --- /dev/null +++ b/python/pyspark/sql/plot/core.py @@ -0,0 +1,135 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Any, TYPE_CHECKING, Optional, Union +from types import ModuleType +from pyspark.errors import PySparkRuntimeError, PySparkValueError +from pyspark.sql.utils import require_minimum_plotly_version + + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + import pandas as pd + from plotly.graph_objs import Figure + + +class PySparkTopNPlotBase: + def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": + from pyspark.sql import SparkSession + + session = SparkSession.getActiveSession() + if session is None: + raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) + + max_rows = int( + session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] + ) + pdf = sdf.limit(max_rows + 1).toPandas() + + self.partial = False + if len(pdf) > max_rows: + self.partial = True + pdf = pdf.iloc[:max_rows] + + return pdf + + +class PySparkSampledPlotBase: + def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": + from pyspark.sql import SparkSession + + session = SparkSession.getActiveSession() + if session is None: + raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) + + sample_ratio = session.conf.get("spark.sql.pyspark.plotting.sample_ratio") + max_rows = int( + session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] + ) + + if sample_ratio is None: + fraction = 1 / (sdf.count() / max_rows) + fraction = min(1.0, fraction) + else: + fraction = float(sample_ratio) + + sampled_sdf = sdf.sample(fraction=fraction) + pdf = sampled_sdf.toPandas() + + return pdf + + +class PySparkPlotAccessor: + plot_data_map = { + "line": PySparkSampledPlotBase().get_sampled, + } + _backends = {} # type: ignore[var-annotated] + + def __init__(self, data: "DataFrame"): + self.data = data + + def __call__( + self, kind: str = "line", backend: Optional[str] = None, **kwargs: Any + ) -> "Figure": + plot_backend = PySparkPlotAccessor._get_plot_backend(backend) + + return plot_backend.plot_pyspark(self.data, kind=kind, **kwargs) + + @staticmethod + def _get_plot_backend(backend: Optional[str] = None) -> ModuleType: + backend = backend or "plotly" + + if backend in PySparkPlotAccessor._backends: + return PySparkPlotAccessor._backends[backend] + + if backend == "plotly": + require_minimum_plotly_version() + else: + raise PySparkValueError( + errorClass="UNSUPPORTED_PLOT_BACKEND", + messageParameters={"backend": backend, "supported_backends": ", ".join(["plotly"])}, + ) + from pyspark.sql.plot import plotly as module + + return module + + def line(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": + """ + Plot DataFrame as lines. + + Parameters + ---------- + x : str + Name of column to use for the horizontal axis. + y : str or list of str + Name(s) of the column(s) to use for the vertical axis. Multiple columns can be plotted. + **kwds : optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.line(x="category", y="int_val") # doctest: +SKIP + >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP + """ + return self(kind="line", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/plot/plotly.py b/python/pyspark/sql/plot/plotly.py new file mode 100644 index 0000000000000..5efc19476057f --- /dev/null +++ b/python/pyspark/sql/plot/plotly.py @@ -0,0 +1,30 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import TYPE_CHECKING, Any + +from pyspark.sql.plot import PySparkPlotAccessor + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + from plotly.graph_objs import Figure + + +def plot_pyspark(data: "DataFrame", kind: str, **kwargs: Any) -> "Figure": + import plotly + + return plotly.plot(PySparkPlotAccessor.plot_data_map[kind](data), kind, **kwargs) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py new file mode 100644 index 0000000000000..c69e438bf7eb0 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py @@ -0,0 +1,36 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.sql.tests.plot.test_frame_plot import DataFramePlotTestsMixin + + +class FramePlotParityTests(DataFramePlotTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_frame_plot import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py new file mode 100644 index 0000000000000..78508fe533379 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py @@ -0,0 +1,36 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.sql.tests.plot.test_frame_plot_plotly import DataFramePlotPlotlyTestsMixin + + +class FramePlotPlotlyParityTests(DataFramePlotPlotlyTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/__init__.py b/python/pyspark/sql/tests/plot/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/sql/tests/plot/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/sql/tests/plot/test_frame_plot.py b/python/pyspark/sql/tests/plot/test_frame_plot.py new file mode 100644 index 0000000000000..19ef53e46b2f4 --- /dev/null +++ b/python/pyspark/sql/tests/plot/test_frame_plot.py @@ -0,0 +1,79 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.errors import PySparkValueError +from pyspark.sql import Row +from pyspark.sql.plot import PySparkSampledPlotBase, PySparkTopNPlotBase +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class DataFramePlotTestsMixin: + def test_backend(self): + accessor = self.spark.range(2).plot + backend = accessor._get_plot_backend() + self.assertEqual(backend.__name__, "pyspark.sql.plot.plotly") + + with self.assertRaises(PySparkValueError) as pe: + accessor._get_plot_backend("matplotlib") + + self.check_error( + exception=pe.exception, + errorClass="UNSUPPORTED_PLOT_BACKEND", + messageParameters={"backend": "matplotlib", "supported_backends": "plotly"}, + ) + + def test_topn_max_rows(self): + try: + self.spark.conf.set("spark.sql.pyspark.plotting.max_rows", "1000") + sdf = self.spark.range(2500) + pdf = PySparkTopNPlotBase().get_top_n(sdf) + self.assertEqual(len(pdf), 1000) + finally: + self.spark.conf.unset("spark.sql.pyspark.plotting.max_rows") + + def test_sampled_plot_with_ratio(self): + try: + self.spark.conf.set("spark.sql.pyspark.plotting.sample_ratio", "0.5") + data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2500)] + sdf = self.spark.createDataFrame(data) + pdf = PySparkSampledPlotBase().get_sampled(sdf) + self.assertEqual(round(len(pdf) / 2500, 1), 0.5) + finally: + self.spark.conf.unset("spark.sql.pyspark.plotting.sample_ratio") + + def test_sampled_plot_with_max_rows(self): + data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2000)] + sdf = self.spark.createDataFrame(data) + pdf = PySparkSampledPlotBase().get_sampled(sdf) + self.assertEqual(round(len(pdf) / 2000, 1), 0.5) + + +class DataFramePlotTests(DataFramePlotTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.plot.test_frame_plot import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py new file mode 100644 index 0000000000000..72a3ed267d192 --- /dev/null +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -0,0 +1,64 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +import pyspark.sql.plot # noqa: F401 +from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message + + +@unittest.skipIf(not have_plotly, plotly_requirement_message) +class DataFramePlotPlotlyTestsMixin: + @property + def sdf(self): + data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + columns = ["category", "int_val", "float_val"] + return self.spark.createDataFrame(data, columns) + + def _check_fig_data(self, fig_data, expected_x, expected_y, expected_name=""): + self.assertEqual(fig_data["mode"], "lines") + self.assertEqual(fig_data["type"], "scatter") + self.assertEqual(fig_data["xaxis"], "x") + self.assertEqual(list(fig_data["x"]), expected_x) + self.assertEqual(fig_data["yaxis"], "y") + self.assertEqual(list(fig_data["y"]), expected_y) + self.assertEqual(fig_data["name"], expected_name) + + def test_line_plot(self): + # single column as vertical axis + fig = self.sdf.plot(kind="line", x="category", y="int_val") + self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20]) + + # multiple columns as vertical axis + fig = self.sdf.plot.line(x="category", y=["int_val", "float_val"]) + self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") + self._check_fig_data(fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + + +class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.plot.test_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py index 11b91612419a3..5d9ec92cbc830 100644 --- a/python/pyspark/sql/utils.py +++ b/python/pyspark/sql/utils.py @@ -41,6 +41,7 @@ PythonException, UnknownException, SparkUpgradeException, + PySparkImportError, PySparkNotImplementedError, PySparkRuntimeError, ) @@ -115,6 +116,22 @@ def require_test_compiled() -> None: ) +def require_minimum_plotly_version() -> None: + """Raise ImportError if plotly is not installed""" + minimum_plotly_version = "4.8" + + try: + import plotly # noqa: F401 + except ImportError as error: + raise PySparkImportError( + errorClass="PACKAGE_NOT_INSTALLED", + messageParameters={ + "package_name": "plotly", + "minimum_version": str(minimum_plotly_version), + }, + ) from error + + class ForeachBatchFunction: """ This is the Python implementation of Java interface 'ForeachBatchFunction'. This wraps diff --git a/python/pyspark/testing/sqlutils.py b/python/pyspark/testing/sqlutils.py index 9f07c44c084cf..00ad40e68bd7c 100644 --- a/python/pyspark/testing/sqlutils.py +++ b/python/pyspark/testing/sqlutils.py @@ -48,6 +48,13 @@ except Exception as e: test_not_compiled_message = str(e) +plotly_requirement_message = None +try: + import plotly +except ImportError as e: + plotly_requirement_message = str(e) +have_plotly = plotly_requirement_message is None + from pyspark.sql import SparkSession from pyspark.sql.types import ArrayType, DoubleType, UserDefinedType, Row from pyspark.testing.utils import ReusedPySparkTestCase, PySparkErrorTestUtils diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index a87b0613292c9..5853e4b66dcc0 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3169,6 +3169,29 @@ object SQLConf { .version("4.0.0") .fallbackConf(Python.PYTHON_WORKER_FAULTHANLDER_ENABLED) + val PYSPARK_PLOT_MAX_ROWS = + buildConf("spark.sql.pyspark.plotting.max_rows") + .doc( + "The visual limit on top-n-based plots. If set to 1000, the first 1000 data points " + + "will be used for plotting.") + .version("4.0.0") + .intConf + .createWithDefault(1000) + + val PYSPARK_PLOT_SAMPLE_RATIO = + buildConf("spark.sql.pyspark.plotting.sample_ratio") + .doc( + "The proportion of data that will be plotted for sample-based plots. It is determined " + + "based on spark.sql.pyspark.plotting.max_rows if not explicitly set." + ) + .version("4.0.0") + .doubleConf + .checkValue( + ratio => ratio >= 0.0 && ratio <= 1.0, + "The value should be between 0.0 and 1.0 inclusive." + ) + .createOptional + val ARROW_SPARKR_EXECUTION_ENABLED = buildConf("spark.sql.execution.arrow.sparkr.enabled") .doc("When true, make use of Apache Arrow for columnar data transfers in SparkR. " + @@ -5855,6 +5878,10 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pythonUDFWorkerFaulthandlerEnabled: Boolean = getConf(PYTHON_UDF_WORKER_FAULTHANLDER_ENABLED) + def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS) + + def pysparkPlotSampleRatio: Option[Double] = getConf(PYSPARK_PLOT_SAMPLE_RATIO) + def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED) def arrowPySparkFallbackEnabled: Boolean = getConf(ARROW_PYSPARK_FALLBACK_ENABLED) From f69b518446e2f18fccdad3e1c23792bbee20f3f5 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Thu, 12 Sep 2024 20:43:33 -0700 Subject: [PATCH 017/189] [SPARK-49620][INFRA] Fix `spark-rm` and `infra` docker files to create `pypy3.9` links ### What changes were proposed in this pull request? This PR aims to fix two Dockerfiles to create `pypy3.9` symlinks instead of `pypy3.8`. https://github.com/apache/spark/blob/d2d293e3fb57d6c9dea084b5fe6707d67c715af3/dev/create-release/spark-rm/Dockerfile#L97 https://github.com/apache/spark/blob/d2d293e3fb57d6c9dea084b5fe6707d67c715af3/dev/infra/Dockerfile#L91 ### Why are the changes needed? Apache Spark 4.0 dropped `Python 3.8` support. We should make it sure that we don't use `pypy3.8` at all. - #46228 ### Does this PR introduce _any_ user-facing change? No. This is a dev-only change. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48095 from dongjoon-hyun/SPARK-49620. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- dev/create-release/spark-rm/Dockerfile | 2 +- dev/infra/Dockerfile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/create-release/spark-rm/Dockerfile b/dev/create-release/spark-rm/Dockerfile index e86b91968bf80..e7f558b523d0c 100644 --- a/dev/create-release/spark-rm/Dockerfile +++ b/dev/create-release/spark-rm/Dockerfile @@ -94,7 +94,7 @@ ENV R_LIBS_SITE "/usr/local/lib/R/site-library:${R_LIBS_SITE}:/usr/lib/R/library RUN add-apt-repository ppa:pypy/ppa RUN mkdir -p /usr/local/pypy/pypy3.9 && \ curl -sqL https://downloads.python.org/pypy/pypy3.9-v7.3.16-linux64.tar.bz2 | tar xjf - -C /usr/local/pypy/pypy3.9 --strip-components=1 && \ - ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3.8 && \ + ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3.9 && \ ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | pypy3 RUN pypy3 -m pip install numpy 'six==1.16.0' 'pandas==2.2.2' scipy coverage matplotlib lxml diff --git a/dev/infra/Dockerfile b/dev/infra/Dockerfile index ce47362999284..5939e429b2f35 100644 --- a/dev/infra/Dockerfile +++ b/dev/infra/Dockerfile @@ -88,7 +88,7 @@ ENV R_LIBS_SITE "/usr/local/lib/R/site-library:${R_LIBS_SITE}:/usr/lib/R/library RUN add-apt-repository ppa:pypy/ppa RUN mkdir -p /usr/local/pypy/pypy3.9 && \ curl -sqL https://downloads.python.org/pypy/pypy3.9-v7.3.16-linux64.tar.bz2 | tar xjf - -C /usr/local/pypy/pypy3.9 --strip-components=1 && \ - ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3.8 && \ + ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3.9 && \ ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | pypy3 RUN pypy3 -m pip install 'numpy==1.26.4' 'six==1.16.0' 'pandas==2.2.2' scipy coverage matplotlib lxml From 23e61f6b1845f0549bc448c7845ee9edae088166 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Thu, 12 Sep 2024 20:48:15 -0700 Subject: [PATCH 018/189] [SPARK-49621][SQL][TESTS] Remove the flaky `EXEC IMMEDIATE STACK OVERFLOW` test case ### What changes were proposed in this pull request? This PR aims to remove the flaky `EXEC IMMEDIATE STACK OVERFLOW` test case. ### Why are the changes needed? To stabilize the CIs. `QueryParsingErrorsSuite` still has a test coverage for the original PR. ### Does this PR introduce _any_ user-facing change? No, this is a test-only change. ### How was this patch tested? Manual review because this is a test case removal. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48096 from dongjoon-hyun/SPARK-49621. Lead-authored-by: Dongjoon Hyun Co-authored-by: YangJie Signed-off-by: Dongjoon Hyun --- .../ExecuteImmediateEndToEndSuite.scala | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExecuteImmediateEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExecuteImmediateEndToEndSuite.scala index 91b1bfd7bf213..62a32da22d957 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExecuteImmediateEndToEndSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExecuteImmediateEndToEndSuite.scala @@ -17,7 +17,6 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{QueryTest} -import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.test.SharedSparkSession class ExecuteImmediateEndToEndSuite extends QueryTest with SharedSparkSession { @@ -37,30 +36,4 @@ class ExecuteImmediateEndToEndSuite extends QueryTest with SharedSparkSession { spark.sql("DROP TEMPORARY VARIABLE IF EXISTS parm;") } } - - test("EXEC IMMEDIATE STACK OVERFLOW") { - try { - spark.sql("DECLARE parm = 1;") - val query = (1 to 20000).map(x => "SELECT 1 as a").mkString(" UNION ALL ") - Seq( - s"EXECUTE IMMEDIATE '$query'", - s"EXECUTE IMMEDIATE '$query' INTO parm").foreach { q => - val e = intercept[ParseException] { - spark.sql(q) - } - - checkError( - exception = e, - condition = "FAILED_TO_PARSE_TOO_COMPLEX", - parameters = Map(), - context = ExpectedContext( - query, - start = 0, - stop = query.length - 1) - ) - } - } finally { - spark.sql("DROP TEMPORARY VARIABLE IF EXISTS parm;") - } - } } From 61814876b26c6fef2dc8238b1aeb0594d9a24472 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Thu, 12 Sep 2024 20:49:16 -0700 Subject: [PATCH 019/189] [SPARK-43354][PYTHON][TESTS] Re-enable `test_create_dataframe_from_pandas_with_day_time_interval` in PyPy3.9 ### What changes were proposed in this pull request? This PR aims to re-enable `test_create_dataframe_from_pandas_with_day_time_interval` in PyPy3.9. ### Why are the changes needed? This was disabled at PyPy3.8, but we dropped Python 3.8 support and the test passed with PyPy3.9. - #46228 **BEFORE: Skipped with `Fails in PyPy Python 3.8, should enable.` message** ``` $ python/run-tests.py --python-executables pypy3 --testnames pyspark.sql.tests.test_creation Running PySpark tests. Output is in /Users/dongjoon/APACHE/spark-merge/python/unit-tests.log Will test against the following Python executables: ['pypy3'] Will test the following Python tests: ['pyspark.sql.tests.test_creation'] pypy3 python_implementation is PyPy pypy3 version is: Python 3.9.19 (a2113ea87262, Apr 21 2024, 05:41:07) [PyPy 7.3.16 with GCC Apple LLVM 15.0.0 (clang-1500.1.0.2.5)] Starting test(pypy3): pyspark.sql.tests.test_creation (temp output: /Users/dongjoon/APACHE/spark-merge/python/target/58e26724-5c3e-4451-80f8-cabdb36f0901/pypy3__pyspark.sql.tests.test_creation__n448ay57.log) Finished test(pypy3): pyspark.sql.tests.test_creation (6s) ... 3 tests were skipped Tests passed in 6 seconds Skipped tests in pyspark.sql.tests.test_creation with pypy3: test_create_dataframe_from_pandas_with_day_time_interval (pyspark.sql.tests.test_creation.DataFrameCreationTests) ... skipped 'Fails in PyPy Python 3.8, should enable.' test_create_dataframe_required_pandas_not_found (pyspark.sql.tests.test_creation.DataFrameCreationTests) ... skipped 'Required Pandas was found.' test_schema_inference_from_pandas_with_dict (pyspark.sql.tests.test_creation.DataFrameCreationTests) ... skipped '[PACKAGE_NOT_INSTALLED] PyArrow >= 10.0.0 must be installed; however, it was not found.' ``` **AFTER** ``` $ python/run-tests.py --python-executables pypy3 --testnames pyspark.sql.tests.test_creation Running PySpark tests. Output is in /Users/dongjoon/APACHE/spark-merge/python/unit-tests.log Will test against the following Python executables: ['pypy3'] Will test the following Python tests: ['pyspark.sql.tests.test_creation'] pypy3 python_implementation is PyPy pypy3 version is: Python 3.9.19 (a2113ea87262, Apr 21 2024, 05:41:07) [PyPy 7.3.16 with GCC Apple LLVM 15.0.0 (clang-1500.1.0.2.5)] Starting test(pypy3): pyspark.sql.tests.test_creation (temp output: /Users/dongjoon/APACHE/spark-merge/python/target/1f0db01f-0beb-4ee2-817f-363eb2f2804d/pypy3__pyspark.sql.tests.test_creation__2w4gy9u1.log) Finished test(pypy3): pyspark.sql.tests.test_creation (13s) ... 2 tests were skipped Tests passed in 13 seconds Skipped tests in pyspark.sql.tests.test_creation with pypy3: test_create_dataframe_required_pandas_not_found (pyspark.sql.tests.test_creation.DataFrameCreationTests) ... skipped 'Required Pandas was found.' test_schema_inference_from_pandas_with_dict (pyspark.sql.tests.test_creation.DataFrameCreationTests) ... skipped '[PACKAGE_NOT_INSTALLED] PyArrow >= 10.0.0 must be installed; however, it was not found.' ``` ### Does this PR introduce _any_ user-facing change? No, this is a test only change. ### How was this patch tested? Manual tests with PyPy3.9. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48097 from dongjoon-hyun/SPARK-43354. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- python/pyspark/sql/tests/test_creation.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/python/pyspark/sql/tests/test_creation.py b/python/pyspark/sql/tests/test_creation.py index dfe66cdd3edf0..c6917aa234b41 100644 --- a/python/pyspark/sql/tests/test_creation.py +++ b/python/pyspark/sql/tests/test_creation.py @@ -15,7 +15,6 @@ # limitations under the License. # -import platform from decimal import Decimal import os import time @@ -111,11 +110,7 @@ def test_create_dataframe_from_pandas_with_dst(self): os.environ["TZ"] = orig_env_tz time.tzset() - # TODO(SPARK-43354): Re-enable test_create_dataframe_from_pandas_with_day_time_interval - @unittest.skipIf( - "pypy" in platform.python_implementation().lower() or not have_pandas, - "Fails in PyPy Python 3.8, should enable.", - ) + @unittest.skipIf(not have_pandas, pandas_requirement_message) # type: ignore def test_create_dataframe_from_pandas_with_day_time_interval(self): # SPARK-37277: Test DayTimeIntervalType in createDataFrame without Arrow. import pandas as pd From e7cf246fb7635ef7b95c18b7958bcadae00aa281 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Thu, 12 Sep 2024 20:52:11 -0700 Subject: [PATCH 020/189] [SPARK-49624][BUILD] Upgrade `aircompressor` to 2.0.2 ### What changes were proposed in this pull request? The pr aims to upgrade `aircompressor` from `0.27` to `2.0.2`. ### Why are the changes needed? https://github.com/airlift/aircompressor/releases/tag/2.0 (ps: 2.0.2 was built against `JDK 1.8`). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48098 from panbingkun/aircompressor_2. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 2 +- pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 2db86ed229a01..e1ac039f25467 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -4,7 +4,7 @@ JTransforms/3.1//JTransforms-3.1.jar RoaringBitmap/1.2.1//RoaringBitmap-1.2.1.jar ST4/4.0.4//ST4-4.0.4.jar activation/1.1.1//activation-1.1.1.jar -aircompressor/0.27//aircompressor-0.27.jar +aircompressor/2.0.2//aircompressor-2.0.2.jar algebra_2.13/2.8.0//algebra_2.13-2.8.0.jar aliyun-java-sdk-core/4.5.10//aliyun-java-sdk-core-4.5.10.jar aliyun-java-sdk-kms/2.11.0//aliyun-java-sdk-kms-2.11.0.jar diff --git a/pom.xml b/pom.xml index b1497c7826855..b9f28eb619258 100644 --- a/pom.xml +++ b/pom.xml @@ -2634,7 +2634,7 @@ io.airlift aircompressor - 0.27 + 2.0.2 org.apache.orc From 319e7cc7d0e7ba9a99f808d51a8d635a6159ce8f Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Fri, 13 Sep 2024 15:00:53 +0800 Subject: [PATCH 021/189] [SPARK-49628][SQL] ConstantFolding should copy stateful expression before evaluating ### What changes were proposed in this pull request? It's possible that a logical plan instance is being shared by multiple DFs and these DFs are executed in parallel. Spark always copy stateful expressions before evaluating them, but one place is missed: `ConstantFolding` can also execute expressions. This PR fixes it. ### Why are the changes needed? avoid concurrency issues. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? Not able to write a test for it, but this concurrency issue is quite obvious ### Was this patch authored or co-authored using generative AI tooling? no Closes #48104 from cloud-fan/constant. Authored-by: Wenchen Fan Signed-off-by: Kent Yao --- .../org/apache/spark/sql/catalyst/optimizer/expressions.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala index d0ee9f2d110d5..3cdde622d51f7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala @@ -79,7 +79,7 @@ object ConstantFolding extends Rule[LogicalPlan] { // Fold expressions that are foldable. case e if e.foldable => try { - Literal.create(e.eval(EmptyRow), e.dataType) + Literal.create(e.freshCopyIfContainsStatefulExpression().eval(EmptyRow), e.dataType) } catch { case NonFatal(_) if isConditionalBranch => // When doing constant folding inside conditional expressions, we should not fail From aa54ed17832f63c177a2dd1b2d0396c4d22adf2e Mon Sep 17 00:00:00 2001 From: panbingkun Date: Fri, 13 Sep 2024 10:58:16 +0200 Subject: [PATCH 022/189] [SPARK-48549][SQL][DOCS][FOLLOWUP] Add migration guide for SQL function `sentences` behavior changes ### What changes were proposed in this pull request? The pr is following up https://github.com/apache/spark/pull/46880, to add migration guide for SQL function `sentences` behavior changes. ### Why are the changes needed? As discussed below: https://github.com/apache/spark/pull/46880#discussion_r1756774840 https://github.com/apache/spark/pull/46880#discussion_r1757102950 All behavior changes need to add migration guide for it. ### Does this PR introduce _any_ user-facing change? Yes, provide clear doc for end-users. ### How was this patch tested? No, only update doc. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48099 from panbingkun/SPARK-48549. Lead-authored-by: panbingkun Co-authored-by: panbingkun Signed-off-by: Max Gekk --- docs/sql-migration-guide.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index ad678c44657ed..0ecd45c2d8c56 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -60,6 +60,7 @@ license: | - Since Spark 4.0, By default views tolerate column type changes in the query and compensate with casts. To restore the previous behavior, allowing up-casts only, set `spark.sql.legacy.viewSchemaCompensation` to `false`. - Since Spark 4.0, Views allow control over how they react to underlying query changes. By default views tolerate column type changes in the query and compensate with casts. To disable this feature set `spark.sql.legacy.viewSchemaBindingMode` to `false`. This also removes the clause from `DESCRIBE EXTENDED` and `SHOW CREATE TABLE`. - Since Spark 4.0, The Storage-Partitioned Join feature flag `spark.sql.sources.v2.bucketing.pushPartValues.enabled` is set to `true`. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.pushPartValues.enabled` to `false`. +- Since Spark 4.0, the `sentences` function uses `Locale(language)` instead of `Locale.US` when `language` parameter is not `NULL` and `country` parameter is `NULL`. ## Upgrading from Spark SQL 3.5.1 to 3.5.2 From 5533c81e34534d43ae90fc2ce5ac1d174d4e8289 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Du=C5=A1an=20Ti=C5=A1ma?= Date: Fri, 13 Sep 2024 15:01:09 +0200 Subject: [PATCH 023/189] [SPARK-48355][SQL] Support for CASE statement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Add support for [case statements](https://docs.google.com/document/d/1cpSuR3KxRuTSJ4ZMQ73FJ4_-hjouNNU2zfI4vri6yhs/edit#heading=h.ofijhkunigv) to sql scripting. There are 2 types of case statement - simple and searched (EXAMPLES BELOW). Proposed changes are: - Add `caseStatement` grammar rule to SqlBaseParser.g4 - Add visit case statement methods to `AstBuilder` - Add `SearchedCaseStatement` and `SearchedCaseStatementExec` classes, to enable them to be run in sql scripts. The reason only searched case nodes are added is that, in the current implementation, a simple case is parsed into a searched case, by creating internal `EqualTo` expressions to compare the main case expression to the expressions in the when clauses. This approach is similar to the existing case **expressions**, which are parsed in the same way. The problem with this approach is that the main expression is unnecessarily evaluated N times, where N is the number of when clauses, which can be quite inefficient, for example if the expression is a complex query. Optimally, the main expression would be evaluated once, and then compared to the other expressions. I'm open to suggestions as to what the best approach to achieve this would be. Simple case compares one expression (case variable) to others, until an equal one is found. Else clause is optional. ``` BEGIN CASE 1 WHEN 1 THEN SELECT 1; WHEN 2 THEN SELECT 2; ELSE SELECT 3; END CASE; END ``` Searched case evaluates boolean expressions. Else clause is optional. ``` BEGIN CASE WHEN 1 = 1 THEN SELECT 1; WHEN 2 IN (1,2,3) THEN SELECT 2; ELSE SELECT 3; END CASE; END ``` ### Why are the changes needed? Case statements are currently not implemented in sql scripting. ### Does this PR introduce _any_ user-facing change? Yes, users will now be able to use case statements in their sql scripts. ### How was this patch tested? Tests for both simple and searched case statements are added to SqlScriptingParserSuite, SqlScriptingExecutionNodeSuite and SqlScriptingInterpreterSuite. ### Was this patch authored or co-authored using generative AI tooling? No Closes #47672 from dusantism-db/sql-scripting-case-statement. Authored-by: Dušan Tišma Signed-off-by: Max Gekk --- .../sql/catalyst/parser/SqlBaseParser.g4 | 8 + .../sql/catalyst/parser/AstBuilder.scala | 48 ++- .../parser/SqlScriptingLogicalOperators.scala | 14 + .../parser/SqlScriptingParserSuite.scala | 297 +++++++++++++- .../scripting/SqlScriptingExecutionNode.scala | 72 ++++ .../scripting/SqlScriptingInterpreter.scala | 13 +- .../SqlScriptingExecutionNodeSuite.scala | 93 +++++ .../SqlScriptingInterpreterSuite.scala | 379 +++++++++++++++++- 8 files changed, 920 insertions(+), 4 deletions(-) diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 42f0094de3515..73d5cb55295ab 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -64,6 +64,7 @@ compoundStatement | setStatementWithOptionalVarKeyword | beginEndCompoundBlock | ifElseStatement + | caseStatement | whileStatement | repeatStatement | leaveStatement @@ -98,6 +99,13 @@ iterateStatement : ITERATE multipartIdentifier ; +caseStatement + : CASE (WHEN conditions+=booleanExpression THEN conditionalBodies+=compoundBody)+ + (ELSE elseBody=compoundBody)? END CASE #searchedCaseStatement + | CASE caseVariable=expression (WHEN conditionExpressions+=expression THEN conditionalBodies+=compoundBody)+ + (ELSE elseBody=compoundBody)? END CASE #simpleCaseStatement + ; + singleStatement : (statement|setResetStatement) SEMICOLON* EOF ; diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index 924b5c2cfeb15..9620ce13d92eb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -261,6 +261,52 @@ class AstBuilder extends DataTypeAstBuilder WhileStatement(condition, body, Some(labelText)) } + override def visitSearchedCaseStatement(ctx: SearchedCaseStatementContext): CaseStatement = { + val conditions = ctx.conditions.asScala.toList.map(boolExpr => withOrigin(boolExpr) { + SingleStatement( + Project( + Seq(Alias(expression(boolExpr), "condition")()), + OneRowRelation())) + }) + val conditionalBodies = + ctx.conditionalBodies.asScala.toList.map(body => visitCompoundBody(body)) + + if (conditions.length != conditionalBodies.length) { + throw SparkException.internalError( + s"Mismatched number of conditions ${conditions.length} and condition bodies" + + s" ${conditionalBodies.length} in case statement") + } + + CaseStatement( + conditions = conditions, + conditionalBodies = conditionalBodies, + elseBody = Option(ctx.elseBody).map(body => visitCompoundBody(body))) + } + + override def visitSimpleCaseStatement(ctx: SimpleCaseStatementContext): CaseStatement = { + // uses EqualTo to compare the case variable(the main case expression) + // to the WHEN clause expressions + val conditions = ctx.conditionExpressions.asScala.toList.map(expr => withOrigin(expr) { + SingleStatement( + Project( + Seq(Alias(EqualTo(expression(ctx.caseVariable), expression(expr)), "condition")()), + OneRowRelation())) + }) + val conditionalBodies = + ctx.conditionalBodies.asScala.toList.map(body => visitCompoundBody(body)) + + if (conditions.length != conditionalBodies.length) { + throw SparkException.internalError( + s"Mismatched number of conditions ${conditions.length} and condition bodies" + + s" ${conditionalBodies.length} in case statement") + } + + CaseStatement( + conditions = conditions, + conditionalBodies = conditionalBodies, + elseBody = Option(ctx.elseBody).map(body => visitCompoundBody(body))) + } + override def visitRepeatStatement(ctx: RepeatStatementContext): RepeatStatement = { val labelText = generateLabelText(Option(ctx.beginLabel()), Option(ctx.endLabel())) val boolExpr = ctx.booleanExpression() @@ -292,7 +338,7 @@ class AstBuilder extends DataTypeAstBuilder case c: RepeatStatementContext if Option(c.beginLabel()).isDefined && c.beginLabel().multipartIdentifier().getText.toLowerCase(Locale.ROOT).equals(label) - => true + => true case _ => false } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala index 5e7e8b0b4fc9a..ed40a5fd734b6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala @@ -124,3 +124,17 @@ case class LeaveStatement(label: String) extends CompoundPlanStatement * @param label Label of the loop to iterate. */ case class IterateStatement(label: String) extends CompoundPlanStatement + +/** + * Logical operator for CASE statement. + * @param conditions Collection of conditions which correspond to WHEN clauses. + * @param conditionalBodies Collection of bodies that have a corresponding condition, + * in WHEN branches. + * @param elseBody Body that is executed if none of the conditions are met, i.e. ELSE branch. + */ +case class CaseStatement( + conditions: Seq[SingleStatement], + conditionalBodies: Seq[CompoundBody], + elseBody: Option[CompoundBody]) extends CompoundPlanStatement { + assert(conditions.length == conditionalBodies.length) +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala index bf527b9c3bd7d..24ad32c5300bc 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala @@ -18,8 +18,9 @@ package org.apache.spark.sql.catalyst.parser import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.{Alias, EqualTo, Expression, In, Literal, ScalarSubquery} import org.apache.spark.sql.catalyst.plans.SQLHelper -import org.apache.spark.sql.catalyst.plans.logical.CreateVariable +import org.apache.spark.sql.catalyst.plans.logical.{CreateVariable, Project} import org.apache.spark.sql.exceptions.SqlScriptingException class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { @@ -1111,6 +1112,287 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { } + test("searched case statement") { + val sqlScriptText = + """ + |BEGIN + | CASE + | WHEN 1 = 1 THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditions.head.getText == "1 = 1") + } + + test("searched case statement - multi when") { + val sqlScriptText = + """ + |BEGIN + | CASE + | WHEN 1 IN (1,2,3) THEN + | SELECT 1; + | WHEN (SELECT * FROM t) THEN + | SELECT * FROM b; + | WHEN 1 = 1 THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 3) + assert(caseStmt.conditionalBodies.length == 3) + assert(caseStmt.elseBody.isEmpty) + + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditions.head.getText == "1 IN (1,2,3)") + + assert(caseStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 1") + + assert(caseStmt.conditions(1).isInstanceOf[SingleStatement]) + assert(caseStmt.conditions(1).getText == "(SELECT * FROM t)") + + assert(caseStmt.conditionalBodies(1).collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies(1).collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT * FROM b") + + assert(caseStmt.conditions(2).isInstanceOf[SingleStatement]) + assert(caseStmt.conditions(2).getText == "1 = 1") + + assert(caseStmt.conditionalBodies(2).collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies(2).collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 42") + } + + test("searched case statement with else") { + val sqlScriptText = + """ + |BEGIN + | CASE + | WHEN 1 = 1 THEN + | SELECT 42; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.elseBody.isDefined) + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditions.head.getText == "1 = 1") + + assert(caseStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 43") + } + + test("searched case statement nested") { + val sqlScriptText = + """ + |BEGIN + | CASE + | WHEN 1 = 1 THEN + | CASE + | WHEN 2 = 1 THEN + | SELECT 41; + | ELSE + | SELECT 42; + | END CASE; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditionalBodies.length == 1) + assert(caseStmt.elseBody.isEmpty) + + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditions.head.getText == "1 = 1") + + assert(caseStmt.conditionalBodies.head.collection.head.isInstanceOf[CaseStatement]) + val nestedCaseStmt = + caseStmt.conditionalBodies.head.collection.head.asInstanceOf[CaseStatement] + + assert(nestedCaseStmt.conditions.length == 1) + assert(nestedCaseStmt.conditionalBodies.length == 1) + assert(nestedCaseStmt.elseBody.isDefined) + + assert(nestedCaseStmt.conditions.head.isInstanceOf[SingleStatement]) + assert(nestedCaseStmt.conditions.head.getText == "2 = 1") + + assert(nestedCaseStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement]) + assert(nestedCaseStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 41") + + assert(nestedCaseStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement]) + assert(nestedCaseStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 42") + } + + test("simple case statement") { + val sqlScriptText = + """ + |BEGIN + | CASE 1 + | WHEN 1 THEN + | SELECT 1; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition(caseStmt.conditions.head, _ == Literal(1), _ == Literal(1)) + } + + + test("simple case statement - multi when") { + val sqlScriptText = + """ + |BEGIN + | CASE 1 + | WHEN 1 THEN + | SELECT 1; + | WHEN (SELECT 2) THEN + | SELECT * FROM b; + | WHEN 3 IN (1,2,3) THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 3) + assert(caseStmt.conditionalBodies.length == 3) + assert(caseStmt.elseBody.isEmpty) + + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition(caseStmt.conditions.head, _ == Literal(1), _ == Literal(1)) + + assert(caseStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 1") + + assert(caseStmt.conditions(1).isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition( + caseStmt.conditions(1), _ == Literal(1), _.isInstanceOf[ScalarSubquery]) + + assert(caseStmt.conditionalBodies(1).collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies(1).collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT * FROM b") + + assert(caseStmt.conditions(2).isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition( + caseStmt.conditions(2), _ == Literal(1), _.isInstanceOf[In]) + + assert(caseStmt.conditionalBodies(2).collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.conditionalBodies(2).collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 42") + } + + test("simple case statement with else") { + val sqlScriptText = + """ + |BEGIN + | CASE 1 + | WHEN 1 THEN + | SELECT 42; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.elseBody.isDefined) + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition(caseStmt.conditions.head, _ == Literal(1), _ == Literal(1)) + + assert(caseStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement]) + assert(caseStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 43") + } + + test("simple case statement nested") { + val sqlScriptText = + """ + |BEGIN + | CASE (SELECT 1) + | WHEN 1 THEN + | CASE 2 + | WHEN 2 THEN + | SELECT 41; + | ELSE + | SELECT 42; + | END CASE; + | END CASE; + |END + |""".stripMargin + val tree = parseScript(sqlScriptText) + assert(tree.collection.length == 1) + assert(tree.collection.head.isInstanceOf[CaseStatement]) + + val caseStmt = tree.collection.head.asInstanceOf[CaseStatement] + assert(caseStmt.conditions.length == 1) + assert(caseStmt.conditionalBodies.length == 1) + assert(caseStmt.elseBody.isEmpty) + + assert(caseStmt.conditions.head.isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition( + caseStmt.conditions.head, _.isInstanceOf[ScalarSubquery], _ == Literal(1)) + + assert(caseStmt.conditionalBodies.head.collection.head.isInstanceOf[CaseStatement]) + val nestedCaseStmt = + caseStmt.conditionalBodies.head.collection.head.asInstanceOf[CaseStatement] + + assert(nestedCaseStmt.conditions.length == 1) + assert(nestedCaseStmt.conditionalBodies.length == 1) + assert(nestedCaseStmt.elseBody.isDefined) + + assert(nestedCaseStmt.conditions.head.isInstanceOf[SingleStatement]) + checkSimpleCaseStatementCondition( + nestedCaseStmt.conditions.head, _ == Literal(2), _ == Literal(2)) + + assert(nestedCaseStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement]) + assert(nestedCaseStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 41") + + assert(nestedCaseStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement]) + assert(nestedCaseStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement] + .getText == "SELECT 42") + } + // Helper methods def cleanupStatementString(statementStr: String): String = { statementStr @@ -1119,4 +1401,17 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { .replace("END", "") .trim } + + private def checkSimpleCaseStatementCondition( + conditionStatement: SingleStatement, + predicateLeft: Expression => Boolean, + predicateRight: Expression => Boolean): Unit = { + assert(conditionStatement.parsedPlan.isInstanceOf[Project]) + val project = conditionStatement.parsedPlan.asInstanceOf[Project] + assert(project.projectList.head.isInstanceOf[Alias]) + assert(project.projectList.head.asInstanceOf[Alias].child.isInstanceOf[EqualTo]) + val equalTo = project.projectList.head.asInstanceOf[Alias].child.asInstanceOf[EqualTo] + assert(predicateLeft(equalTo.left)) + assert(predicateRight(equalTo.right)) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala index cae7976143142..af9fd5464277c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala @@ -405,6 +405,78 @@ class WhileStatementExec( } } +/** + * Executable node for CaseStatement. + * @param conditions Collection of executable conditions which correspond to WHEN clauses. + * @param conditionalBodies Collection of executable bodies that have a corresponding condition, + * in WHEN branches. + * @param elseBody Body that is executed if none of the conditions are met, i.e. ELSE branch. + * @param session Spark session that SQL script is executed within. + */ +class CaseStatementExec( + conditions: Seq[SingleStatementExec], + conditionalBodies: Seq[CompoundBodyExec], + elseBody: Option[CompoundBodyExec], + session: SparkSession) extends NonLeafStatementExec { + private object CaseState extends Enumeration { + val Condition, Body = Value + } + + private var state = CaseState.Condition + private var curr: Option[CompoundStatementExec] = Some(conditions.head) + + private var clauseIdx: Int = 0 + private val conditionsCount = conditions.length + + private lazy val treeIterator: Iterator[CompoundStatementExec] = + new Iterator[CompoundStatementExec] { + override def hasNext: Boolean = curr.nonEmpty + + override def next(): CompoundStatementExec = state match { + case CaseState.Condition => + val condition = curr.get.asInstanceOf[SingleStatementExec] + if (evaluateBooleanCondition(session, condition)) { + state = CaseState.Body + curr = Some(conditionalBodies(clauseIdx)) + } else { + clauseIdx += 1 + if (clauseIdx < conditionsCount) { + // There are WHEN clauses remaining. + state = CaseState.Condition + curr = Some(conditions(clauseIdx)) + } else if (elseBody.isDefined) { + // ELSE clause exists. + state = CaseState.Body + curr = Some(elseBody.get) + } else { + // No remaining clauses. + curr = None + } + } + condition + case CaseState.Body => + assert(curr.get.isInstanceOf[CompoundBodyExec]) + val currBody = curr.get.asInstanceOf[CompoundBodyExec] + val retStmt = currBody.getTreeIterator.next() + if (!currBody.getTreeIterator.hasNext) { + curr = None + } + retStmt + } + } + + override def getTreeIterator: Iterator[CompoundStatementExec] = treeIterator + + override def reset(): Unit = { + state = CaseState.Condition + curr = Some(conditions.head) + clauseIdx = 0 + conditions.foreach(c => c.reset()) + conditionalBodies.foreach(b => b.reset()) + elseBody.foreach(b => b.reset()) + } +} + /** * Executable node for RepeatStatement. * @param condition Executable node for the condition - evaluates to a row with a single boolean diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala index 865b33999655a..917b4d6f45ee0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.scripting import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.analysis.UnresolvedIdentifier -import org.apache.spark.sql.catalyst.parser.{CompoundBody, CompoundPlanStatement, IfElseStatement, IterateStatement, LeaveStatement, RepeatStatement, SingleStatement, WhileStatement} +import org.apache.spark.sql.catalyst.parser.{CaseStatement, CompoundBody, CompoundPlanStatement, IfElseStatement, IterateStatement, LeaveStatement, RepeatStatement, SingleStatement, WhileStatement} import org.apache.spark.sql.catalyst.plans.logical.{CreateVariable, DropVariable, LogicalPlan} import org.apache.spark.sql.catalyst.trees.Origin @@ -95,6 +95,17 @@ case class SqlScriptingInterpreter() { new IfElseStatementExec( conditionsExec, conditionalBodiesExec, unconditionalBodiesExec, session) + case CaseStatement(conditions, conditionalBodies, elseBody) => + val conditionsExec = conditions.map(condition => + // todo: what to put here for isInternal, in case of simple case statement + new SingleStatementExec(condition.parsedPlan, condition.origin, isInternal = false)) + val conditionalBodiesExec = conditionalBodies.map(body => + transformTreeIntoExecutable(body, session).asInstanceOf[CompoundBodyExec]) + val unconditionalBodiesExec = elseBody.map(body => + transformTreeIntoExecutable(body, session).asInstanceOf[CompoundBodyExec]) + new CaseStatementExec( + conditionsExec, conditionalBodiesExec, unconditionalBodiesExec, session) + case WhileStatement(condition, body, label) => val conditionExec = new SingleStatementExec(condition.parsedPlan, condition.origin, isInternal = false) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala index 4b72ca8ecaa97..83d8191d01ec1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala @@ -576,4 +576,97 @@ class SqlScriptingExecutionNodeSuite extends SparkFunSuite with SharedSparkSessi "body1", "lbl", "con1", "body1", "lbl", "con1")) } + + test("searched case - enter first WHEN clause") { + val iter = new CompoundBodyExec(Seq( + new CaseStatementExec( + conditions = Seq( + TestIfElseCondition(condVal = true, description = "con1"), + TestIfElseCondition(condVal = false, description = "con2") + ), + conditionalBodies = Seq( + new CompoundBodyExec(Seq(TestLeafStatement("body1"))), + new CompoundBodyExec(Seq(TestLeafStatement("body2"))) + ), + elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body3")))), + session = spark + ) + )).getTreeIterator + val statements = iter.map(extractStatementValue).toSeq + assert(statements === Seq("con1", "body1")) + } + + test("searched case - enter body of the ELSE clause") { + val iter = new CompoundBodyExec(Seq( + new CaseStatementExec( + conditions = Seq( + TestIfElseCondition(condVal = false, description = "con1") + ), + conditionalBodies = Seq( + new CompoundBodyExec(Seq(TestLeafStatement("body1"))) + ), + elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body2")))), + session = spark + ) + )).getTreeIterator + val statements = iter.map(extractStatementValue).toSeq + assert(statements === Seq("con1", "body2")) + } + + test("searched case - enter second WHEN clause") { + val iter = new CompoundBodyExec(Seq( + new CaseStatementExec( + conditions = Seq( + TestIfElseCondition(condVal = false, description = "con1"), + TestIfElseCondition(condVal = true, description = "con2") + ), + conditionalBodies = Seq( + new CompoundBodyExec(Seq(TestLeafStatement("body1"))), + new CompoundBodyExec(Seq(TestLeafStatement("body2"))) + ), + elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body3")))), + session = spark + ) + )).getTreeIterator + val statements = iter.map(extractStatementValue).toSeq + assert(statements === Seq("con1", "con2", "body2")) + } + + test("searched case - without else (successful check)") { + val iter = new CompoundBodyExec(Seq( + new CaseStatementExec( + conditions = Seq( + TestIfElseCondition(condVal = false, description = "con1"), + TestIfElseCondition(condVal = true, description = "con2") + ), + conditionalBodies = Seq( + new CompoundBodyExec(Seq(TestLeafStatement("body1"))), + new CompoundBodyExec(Seq(TestLeafStatement("body2"))) + ), + elseBody = None, + session = spark + ) + )).getTreeIterator + val statements = iter.map(extractStatementValue).toSeq + assert(statements === Seq("con1", "con2", "body2")) + } + + test("searched case - without else (unsuccessful checks)") { + val iter = new CompoundBodyExec(Seq( + new CaseStatementExec( + conditions = Seq( + TestIfElseCondition(condVal = false, description = "con1"), + TestIfElseCondition(condVal = false, description = "con2") + ), + conditionalBodies = Seq( + new CompoundBodyExec(Seq(TestLeafStatement("body1"))), + new CompoundBodyExec(Seq(TestLeafStatement("body2"))) + ), + elseBody = None, + session = spark + ) + )).getTreeIterator + val statements = iter.map(extractStatementValue).toSeq + assert(statements === Seq("con1", "con2")) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala index 8d9cd1d8c780e..4851faf897a02 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.scripting -import org.apache.spark.SparkException +import org.apache.spark.{SparkException, SparkNumberFormatException} import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, QueryTest, Row} import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.exceptions.SqlScriptingException @@ -368,6 +368,383 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { } } + test("searched case") { + val commands = + """ + |BEGIN + | CASE + | WHEN 1 = 1 THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("searched case nested") { + val commands = + """ + |BEGIN + | CASE + | WHEN 1=1 THEN + | CASE + | WHEN 2=1 THEN + | SELECT 41; + | ELSE + | SELECT 42; + | END CASE; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("searched case second case") { + val commands = + """ + |BEGIN + | CASE + | WHEN 1 = (SELECT 2) THEN + | SELECT 1; + | WHEN 2 = 2 THEN + | SELECT 42; + | WHEN (SELECT * FROM t) THEN + | SELECT * FROM b; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("searched case going in else") { + val commands = + """ + |BEGIN + | CASE + | WHEN 2 = 1 THEN + | SELECT 1; + | WHEN 3 IN (1,2) THEN + | SELECT 2; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(43))) + verifySqlScriptResult(commands, expected) + } + + test("searched case with count") { + withTable("t") { + val commands = + """ + |BEGIN + |CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet; + |INSERT INTO t VALUES (1, 'a', 1.0); + |INSERT INTO t VALUES (1, 'a', 1.0); + |CASE + | WHEN (SELECT COUNT(*) > 2 FROM t) THEN + | SELECT 42; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + + val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(43))) + verifySqlScriptResult(commands, expected) + } + } + + test("searched case else with count") { + withTable("t") { + val commands = + """ + |BEGIN + | CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet; + | INSERT INTO t VALUES (1, 'a', 1.0); + | INSERT INTO t VALUES (1, 'a', 1.0); + | CASE + | WHEN (SELECT COUNT(*) > 2 FROM t) THEN + | SELECT 42; + | WHEN (SELECT COUNT(*) > 1 FROM t) THEN + | SELECT 43; + | ELSE + | SELECT 44; + | END CASE; + |END + |""".stripMargin + + val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(43))) + verifySqlScriptResult(commands, expected) + } + } + + test("searched case no cases matched no else") { + val commands = + """ + |BEGIN + | CASE + | WHEN 1 = 2 THEN + | SELECT 42; + | WHEN 1 = 3 THEN + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val expected = Seq() + verifySqlScriptResult(commands, expected) + } + + test("searched case when evaluates to null") { + withTable("t") { + val commands = + """ + |BEGIN + | CREATE TABLE t (a BOOLEAN) USING parquet; + | CASE + | WHEN (SELECT * FROM t) THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + + checkError( + exception = intercept[SqlScriptingException] ( + runSqlScript(commands) + ), + condition = "BOOLEAN_STATEMENT_WITH_EMPTY_ROW", + parameters = Map("invalidStatement" -> "(SELECT * FROM T)") + ) + } + } + + test("searched case with non boolean condition - constant") { + val commands = + """ + |BEGIN + | CASE + | WHEN 1 THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + + checkError( + exception = intercept[SqlScriptingException] ( + runSqlScript(commands) + ), + condition = "INVALID_BOOLEAN_STATEMENT", + parameters = Map("invalidStatement" -> "1") + ) + } + + test("searched case with too many rows in subquery condition") { + withTable("t") { + val commands = + """ + |BEGIN + | CREATE TABLE t (a BOOLEAN) USING parquet; + | INSERT INTO t VALUES (true); + | INSERT INTO t VALUES (true); + | CASE + | WHEN (SELECT * FROM t) THEN + | SELECT 1; + | END CASE; + |END + |""".stripMargin + + checkError( + exception = intercept[SparkException] ( + runSqlScript(commands) + ), + condition = "SCALAR_SUBQUERY_TOO_MANY_ROWS", + parameters = Map.empty, + context = ExpectedContext(fragment = "(SELECT * FROM t)", start = 124, stop = 140) + ) + } + } + + test("simple case") { + val commands = + """ + |BEGIN + | CASE 1 + | WHEN 1 THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("simple case nested") { + val commands = + """ + |BEGIN + | CASE 1 + | WHEN 1 THEN + | CASE 2 + | WHEN (SELECT 3) THEN + | SELECT 41; + | ELSE + | SELECT 42; + | END CASE; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("simple case second case") { + val commands = + """ + |BEGIN + | CASE (SELECT 2) + | WHEN 1 THEN + | SELECT 1; + | WHEN 2 THEN + | SELECT 42; + | WHEN (SELECT * FROM t) THEN + | SELECT * FROM b; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + + test("simple case going in else") { + val commands = + """ + |BEGIN + | CASE 1 + | WHEN 2 THEN + | SELECT 1; + | WHEN 3 THEN + | SELECT 2; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val expected = Seq(Seq(Row(43))) + verifySqlScriptResult(commands, expected) + } + + test("simple case with count") { + withTable("t") { + val commands = + """ + |BEGIN + |CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet; + |INSERT INTO t VALUES (1, 'a', 1.0); + |INSERT INTO t VALUES (1, 'a', 1.0); + |CASE (SELECT COUNT(*) FROM t) + | WHEN 1 THEN + | SELECT 41; + | WHEN 2 THEN + | SELECT 42; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + + val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(42))) + verifySqlScriptResult(commands, expected) + } + } + + test("simple case else with count") { + withTable("t") { + val commands = + """ + |BEGIN + | CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet; + | INSERT INTO t VALUES (1, 'a', 1.0); + | INSERT INTO t VALUES (2, 'b', 2.0); + | CASE (SELECT COUNT(*) FROM t) + | WHEN 1 THEN + | SELECT 42; + | WHEN 3 THEN + | SELECT 43; + | ELSE + | SELECT 44; + | END CASE; + |END + |""".stripMargin + + val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(44))) + verifySqlScriptResult(commands, expected) + } + } + + test("simple case no cases matched no else") { + val commands = + """ + |BEGIN + | CASE 1 + | WHEN 2 THEN + | SELECT 42; + | WHEN 3 THEN + | SELECT 43; + | END CASE; + |END + |""".stripMargin + val expected = Seq() + verifySqlScriptResult(commands, expected) + } + + test("simple case mismatched types") { + val commands = + """ + |BEGIN + | CASE 1 + | WHEN "one" THEN + | SELECT 42; + | END CASE; + |END + |""".stripMargin + + checkError( + exception = intercept[SparkNumberFormatException] ( + runSqlScript(commands) + ), + condition = "CAST_INVALID_INPUT", + parameters = Map( + "expression" -> "'one'", + "sourceType" -> "\"STRING\"", + "targetType" -> "\"BIGINT\""), + context = ExpectedContext(fragment = "\"one\"", start = 23, stop = 27) + ) + } + + test("simple case compare with null") { + withTable("t") { + val commands = + """ + |BEGIN + | CREATE TABLE t (a INT) USING parquet; + | CASE (SELECT COUNT(*) FROM t) + | WHEN 1 THEN + | SELECT 42; + | ELSE + | SELECT 43; + | END CASE; + |END + |""".stripMargin + + val expected = Seq(Seq.empty[Row], Seq(Row(43))) + verifySqlScriptResult(commands, expected) + } + } + test("if's condition must be a boolean statement") { withTable("t") { val commands = From 9fc58aa4c0753ef42e0172d73e13b45a00a730f9 Mon Sep 17 00:00:00 2001 From: beliefer Date: Fri, 13 Sep 2024 22:18:05 +0800 Subject: [PATCH 024/189] [SPARK-49488][SQL] MySQL dialect supports pushdown datetime functions ### What changes were proposed in this pull request? This PR propose to make MySQL dialect supports pushdown datetime functions. ### Why are the changes needed? Currently, DS V2 pushdown framework pushed the datetime functions with in a common way. But MySQL doesn't support some datetime functions. ### Does this PR introduce _any_ user-facing change? 'No'. This is a new feature for MySQL dialect. ### How was this patch tested? GA. ### Was this patch authored or co-authored using generative AI tooling? 'No'. Closes #47951 from beliefer/SPARK-49488. Authored-by: beliefer Signed-off-by: Wenchen Fan --- .../sql/jdbc/v2/MySQLIntegrationSuite.scala | 86 ++++++++++++++++++- .../apache/spark/sql/jdbc/v2/V2JDBCTest.scala | 8 +- .../apache/spark/sql/jdbc/MySQLDialect.scala | 23 ++++- 3 files changed, 114 insertions(+), 3 deletions(-) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala index 60685f5c0c6b9..700c05b54a256 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala @@ -77,8 +77,19 @@ class MySQLIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTest s"""CREATE TABLE pattern_testing_table ( |pattern_testing_col LONGTEXT |) - """.stripMargin + |""".stripMargin ).executeUpdate() + connection.prepareStatement( + "CREATE TABLE datetime (name VARCHAR(32), date1 DATE, time1 TIMESTAMP)") + .executeUpdate() + } + + override def dataPreparation(connection: Connection): Unit = { + super.dataPreparation(connection) + connection.prepareStatement("INSERT INTO datetime VALUES " + + "('amy', '2022-05-19', '2022-05-19 00:00:00')").executeUpdate() + connection.prepareStatement("INSERT INTO datetime VALUES " + + "('alex', '2022-05-18', '2022-05-18 00:00:00')").executeUpdate() } override def testUpdateColumnType(tbl: String): Unit = { @@ -157,6 +168,79 @@ class MySQLIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTest assert(sql(s"SELECT char_length(c1) from $tableName").head().get(0) === 65536) } } + + override def testDatetime(tbl: String): Unit = { + val df1 = sql(s"SELECT name FROM $tbl WHERE " + + "dayofyear(date1) > 100 AND dayofmonth(date1) > 10 ") + checkFilterPushed(df1) + val rows1 = df1.collect() + assert(rows1.length === 2) + assert(rows1(0).getString(0) === "amy") + assert(rows1(1).getString(0) === "alex") + + val df2 = sql(s"SELECT name FROM $tbl WHERE year(date1) = 2022 AND quarter(date1) = 2") + checkFilterPushed(df2) + val rows2 = df2.collect() + assert(rows2.length === 2) + assert(rows2(0).getString(0) === "amy") + assert(rows2(1).getString(0) === "alex") + + val df3 = sql(s"SELECT name FROM $tbl WHERE second(time1) = 0 AND month(date1) = 5") + checkFilterPushed(df3) + val rows3 = df3.collect() + assert(rows3.length === 2) + assert(rows3(0).getString(0) === "amy") + assert(rows3(1).getString(0) === "alex") + + val df4 = sql(s"SELECT name FROM $tbl WHERE hour(time1) = 0 AND minute(time1) = 0") + checkFilterPushed(df4) + val rows4 = df4.collect() + assert(rows4.length === 2) + assert(rows4(0).getString(0) === "amy") + assert(rows4(1).getString(0) === "alex") + + val df5 = sql(s"SELECT name FROM $tbl WHERE " + + "extract(WEEk from date1) > 10 AND extract(YEAROFWEEK from date1) = 2022") + checkFilterPushed(df5) + val rows5 = df5.collect() + assert(rows5.length === 2) + assert(rows5(0).getString(0) === "amy") + assert(rows5(1).getString(0) === "alex") + + val df6 = sql(s"SELECT name FROM $tbl WHERE date_add(date1, 1) = date'2022-05-20' " + + "AND datediff(date1, '2022-05-10') > 0") + checkFilterPushed(df6) + val rows6 = df6.collect() + assert(rows6.length === 1) + assert(rows6(0).getString(0) === "amy") + + val df7 = sql(s"SELECT name FROM $tbl WHERE weekday(date1) = 2") + checkFilterPushed(df7) + val rows7 = df7.collect() + assert(rows7.length === 1) + assert(rows7(0).getString(0) === "alex") + + val df8 = sql(s"SELECT name FROM $tbl WHERE dayofweek(date1) = 4") + checkFilterPushed(df8) + val rows8 = df8.collect() + assert(rows8.length === 1) + assert(rows8(0).getString(0) === "alex") + + val df9 = sql(s"SELECT name FROM $tbl WHERE " + + "dayofyear(date1) > 100 order by dayofyear(date1) limit 1") + checkFilterPushed(df9) + val rows9 = df9.collect() + assert(rows9.length === 1) + assert(rows9(0).getString(0) === "alex") + + // MySQL does not support + val df10 = sql(s"SELECT name FROM $tbl WHERE trunc(date1, 'week') = date'2022-05-16'") + checkFilterPushed(df10, false) + val rows10 = df10.collect() + assert(rows10.length === 2) + assert(rows10(0).getString(0) === "amy") + assert(rows10(1).getString(0) === "alex") + } } /** diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala index d3629d871cd42..54635f69f8b65 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala @@ -353,7 +353,7 @@ private[v2] trait V2JDBCTest extends SharedSparkSession with DockerIntegrationFu } } - private def checkFilterPushed(df: DataFrame, pushed: Boolean = true): Unit = { + protected def checkFilterPushed(df: DataFrame, pushed: Boolean = true): Unit = { val filter = df.queryExecution.optimizedPlan.collect { case f: Filter => f } @@ -980,4 +980,10 @@ private[v2] trait V2JDBCTest extends SharedSparkSession with DockerIntegrationFu ) } } + + def testDatetime(tbl: String): Unit = {} + + test("scan with filter push-down with date time functions") { + testDatetime(s"$catalogAndNamespace.${caseConvert("datetime")}") + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala index f2b626490d13c..785bf5b13aa78 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala @@ -46,12 +46,33 @@ private case class MySQLDialect() extends JdbcDialect with SQLConfHelper with No // See https://dev.mysql.com/doc/refman/8.0/en/aggregate-functions.html private val supportedAggregateFunctions = Set("MAX", "MIN", "SUM", "COUNT", "AVG") ++ distinctUnsupportedAggregateFunctions - private val supportedFunctions = supportedAggregateFunctions + private val supportedFunctions = supportedAggregateFunctions ++ Set("DATE_ADD", "DATE_DIFF") override def isSupportedFunction(funcName: String): Boolean = supportedFunctions.contains(funcName) class MySQLSQLBuilder extends JDBCSQLBuilder { + override def visitExtract(field: String, source: String): String = { + field match { + case "DAY_OF_YEAR" => s"DAYOFYEAR($source)" + case "YEAR_OF_WEEK" => s"EXTRACT(YEAR FROM $source)" + // WEEKDAY uses Monday = 0, Tuesday = 1, ... and ISO standard is Monday = 1, ..., + // so we use the formula (WEEKDAY + 1) to follow the ISO standard. + case "DAY_OF_WEEK" => s"(WEEKDAY($source) + 1)" + case _ => super.visitExtract(field, source) + } + } + + override def visitSQLFunction(funcName: String, inputs: Array[String]): String = { + funcName match { + case "DATE_ADD" => + s"DATE_ADD(${inputs(0)}, INTERVAL ${inputs(1)} DAY)" + case "DATE_DIFF" => + s"DATEDIFF(${inputs(0)}, ${inputs(1)})" + case _ => super.visitSQLFunction(funcName, inputs) + } + } + override def visitSortOrder( sortKey: String, sortDirection: SortDirection, nullOrdering: NullOrdering): String = { (sortDirection, nullOrdering) match { From f92e9489fb23a85195067cd0f0f5cd9e9d00b138 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Fri, 13 Sep 2024 09:45:46 -0700 Subject: [PATCH 025/189] [SPARK-49234][BUILD][FOLLOWUP] Add `LICENSE-xz.txt` to `licenses-binary` folder ### What changes were proposed in this pull request? This PR aims to add `LICENSE-xz.txt` to `licenses-binary` folder. ### Why are the changes needed? To provide the license properly. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Manual review. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48107 from dongjoon-hyun/SPARK-49234-2. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- licenses-binary/LICENSE-xz.txt | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 licenses-binary/LICENSE-xz.txt diff --git a/licenses-binary/LICENSE-xz.txt b/licenses-binary/LICENSE-xz.txt new file mode 100644 index 0000000000000..4322122aecf1a --- /dev/null +++ b/licenses-binary/LICENSE-xz.txt @@ -0,0 +1,11 @@ +Permission to use, copy, modify, and/or distribute this +software for any purpose with or without fee is hereby granted. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL +WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL +THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR +CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, +NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN +CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. From b9b64fe37e93055674f1e0796f8dc40ec8b40992 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Du=C5=A1an=20Ti=C5=A1ma?= Date: Fri, 13 Sep 2024 19:01:03 +0200 Subject: [PATCH 026/189] [SPARK-49244][SQL] Further exception improvements for parser/interpreter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Improved SQL scripting exceptions, by removing duplicate line numbers from error messages and adding backquotes to label identifiers. Additionally, expanded exception tests so they also check if the correct line numbers are being shown. ### Why are the changes needed? They are needed in order to make error messages for sql scripting more readable. ### Does this PR introduce _any_ user-facing change? Yes, sql scripting error messages will now be more readable. ### How was this patch tested? Existing exception tests were improved to check for correct line numbers. ### Was this patch authored or co-authored using generative AI tooling? No Closes #47803 from dusantism-db/sql-scripting-further-exception-improvements. Authored-by: Dušan Tišma Signed-off-by: Max Gekk --- .../resources/error/error-conditions.json | 4 +- .../sql/catalyst/parser/AstBuilder.scala | 12 +++--- .../spark/sql/errors/SqlScriptingErrors.scala | 15 +++---- .../exceptions/SqlScriptingException.scala | 2 +- .../parser/SqlScriptingParserSuite.scala | 43 +++++++++++-------- .../SqlScriptingInterpreterSuite.scala | 18 +++++--- 6 files changed, 52 insertions(+), 42 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 0ebeea9aed8d2..229da4fa17de7 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -3118,12 +3118,12 @@ "subClass" : { "NOT_ALLOWED_IN_SCOPE" : { "message" : [ - "Variable was declared on line , which is not allowed in this scope." + "Declaration of the variable is not allowed in this scope." ] }, "ONLY_AT_BEGINNING" : { "message" : [ - "Variable can only be declared at the beginning of the compound, but it was declared on line ." + "Variable can only be declared at the beginning of the compound." ] } }, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index 9620ce13d92eb..7ad7d60e70c96 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -173,14 +173,10 @@ class AstBuilder extends DataTypeAstBuilder case Some(c: CreateVariable) => if (allowVarDeclare) { throw SqlScriptingErrors.variableDeclarationOnlyAtBeginning( - c.origin, - toSQLId(c.name.asInstanceOf[UnresolvedIdentifier].nameParts), - c.origin.line.get.toString) + c.origin, c.name.asInstanceOf[UnresolvedIdentifier].nameParts) } else { throw SqlScriptingErrors.variableDeclarationNotAllowedInScope( - c.origin, - toSQLId(c.name.asInstanceOf[UnresolvedIdentifier].nameParts), - c.origin.line.get.toString) + c.origin, c.name.asInstanceOf[UnresolvedIdentifier].nameParts) } case _ => } @@ -200,7 +196,9 @@ class AstBuilder extends DataTypeAstBuilder el.multipartIdentifier().getText.toLowerCase(Locale.ROOT) => withOrigin(bl) { throw SqlScriptingErrors.labelsMismatch( - CurrentOrigin.get, bl.multipartIdentifier().getText, el.multipartIdentifier().getText) + CurrentOrigin.get, + bl.multipartIdentifier().getText, + el.multipartIdentifier().getText) } case (None, Some(el: EndLabelContext)) => withOrigin(el) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala index 591d2e3e53d47..7f13dc334e06e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.errors import org.apache.spark.sql.catalyst.trees.Origin +import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLStmt import org.apache.spark.sql.exceptions.SqlScriptingException @@ -32,7 +33,7 @@ private[sql] object SqlScriptingErrors { origin = origin, errorClass = "LABELS_MISMATCH", cause = null, - messageParameters = Map("beginLabel" -> beginLabel, "endLabel" -> endLabel)) + messageParameters = Map("beginLabel" -> toSQLId(beginLabel), "endLabel" -> toSQLId(endLabel))) } def endLabelWithoutBeginLabel(origin: Origin, endLabel: String): Throwable = { @@ -40,29 +41,27 @@ private[sql] object SqlScriptingErrors { origin = origin, errorClass = "END_LABEL_WITHOUT_BEGIN_LABEL", cause = null, - messageParameters = Map("endLabel" -> endLabel)) + messageParameters = Map("endLabel" -> toSQLId(endLabel))) } def variableDeclarationNotAllowedInScope( origin: Origin, - varName: String, - lineNumber: String): Throwable = { + varName: Seq[String]): Throwable = { new SqlScriptingException( origin = origin, errorClass = "INVALID_VARIABLE_DECLARATION.NOT_ALLOWED_IN_SCOPE", cause = null, - messageParameters = Map("varName" -> varName, "lineNumber" -> lineNumber)) + messageParameters = Map("varName" -> toSQLId(varName))) } def variableDeclarationOnlyAtBeginning( origin: Origin, - varName: String, - lineNumber: String): Throwable = { + varName: Seq[String]): Throwable = { new SqlScriptingException( origin = origin, errorClass = "INVALID_VARIABLE_DECLARATION.ONLY_AT_BEGINNING", cause = null, - messageParameters = Map("varName" -> varName, "lineNumber" -> lineNumber)) + messageParameters = Map("varName" -> toSQLId(varName))) } def invalidBooleanStatement( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/exceptions/SqlScriptingException.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/exceptions/SqlScriptingException.scala index 4354e7e3635e4..f0c28c95046eb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/exceptions/SqlScriptingException.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/exceptions/SqlScriptingException.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.exceptions.SqlScriptingException.errorMessageWithLin class SqlScriptingException ( errorClass: String, cause: Throwable, - origin: Origin, + val origin: Origin, messageParameters: Map[String, String] = Map.empty) extends Exception( errorMessageWithLineNumber(Option(origin), errorClass, messageParameters), diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala index 24ad32c5300bc..ba634333e06fb 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala @@ -21,6 +21,7 @@ import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.expressions.{Alias, EqualTo, Expression, In, Literal, ScalarSubquery} import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.plans.logical.{CreateVariable, Project} +import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.exceptions.SqlScriptingException class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { @@ -206,13 +207,14 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { | SELECT a, b, c FROM T; | SELECT * FROM T; |END lbl_end""".stripMargin - + val exception = intercept[SqlScriptingException] { + parseScript(sqlScriptText) + } checkError( - exception = intercept[SqlScriptingException] { - parseScript(sqlScriptText) - }, + exception = exception, condition = "LABELS_MISMATCH", - parameters = Map("beginLabel" -> "lbl_begin", "endLabel" -> "lbl_end")) + parameters = Map("beginLabel" -> toSQLId("lbl_begin"), "endLabel" -> toSQLId("lbl_end"))) + assert(exception.origin.line.contains(2)) } test("compound: endLabel") { @@ -225,13 +227,14 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { | SELECT a, b, c FROM T; | SELECT * FROM T; |END lbl""".stripMargin - + val exception = intercept[SqlScriptingException] { + parseScript(sqlScriptText) + } checkError( - exception = intercept[SqlScriptingException] { - parseScript(sqlScriptText) - }, + exception = exception, condition = "END_LABEL_WITHOUT_BEGIN_LABEL", - parameters = Map("endLabel" -> "lbl")) + parameters = Map("endLabel" -> toSQLId("lbl"))) + assert(exception.origin.line.contains(8)) } test("compound: beginLabel + endLabel with different casing") { @@ -287,12 +290,14 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { | SELECT 1; | DECLARE testVariable INTEGER; |END""".stripMargin + val exception = intercept[SqlScriptingException] { + parseScript(sqlScriptText) + } checkError( - exception = intercept[SqlScriptingException] { - parseScript(sqlScriptText) - }, + exception = exception, condition = "INVALID_VARIABLE_DECLARATION.ONLY_AT_BEGINNING", - parameters = Map("varName" -> "`testVariable`", "lineNumber" -> "4")) + parameters = Map("varName" -> "`testVariable`")) + assert(exception.origin.line.contains(4)) } test("declare in wrong scope") { @@ -303,12 +308,14 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper { | DECLARE testVariable INTEGER; | END IF; |END""".stripMargin + val exception = intercept[SqlScriptingException] { + parseScript(sqlScriptText) + } checkError( - exception = intercept[SqlScriptingException] { - parseScript(sqlScriptText) - }, + exception = exception, condition = "INVALID_VARIABLE_DECLARATION.NOT_ALLOWED_IN_SCOPE", - parameters = Map("varName" -> "`testVariable`", "lineNumber" -> "4")) + parameters = Map("varName" -> "`testVariable`")) + assert(exception.origin.line.contains(4)) } test("SET VAR statement test") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala index 4851faf897a02..3fad99eba509a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala @@ -755,13 +755,16 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { | END IF; |END |""".stripMargin + val exception = intercept[SqlScriptingException] { + runSqlScript(commands) + } checkError( - exception = intercept[SqlScriptingException] ( - runSqlScript(commands) - ), + exception = exception, condition = "INVALID_BOOLEAN_STATEMENT", parameters = Map("invalidStatement" -> "1") ) + assert(exception.origin.line.isDefined) + assert(exception.origin.line.get == 3) } } @@ -777,13 +780,16 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { | END IF; |END |""".stripMargin + val exception = intercept[SqlScriptingException] { + runSqlScript(commands1) + } checkError( - exception = intercept[SqlScriptingException] ( - runSqlScript(commands1) - ), + exception = exception, condition = "BOOLEAN_STATEMENT_WITH_EMPTY_ROW", parameters = Map("invalidStatement" -> "(SELECT * FROM T1)") ) + assert(exception.origin.line.isDefined) + assert(exception.origin.line.get == 4) // too many rows ( > 1 ) val commands2 = From bbe6d573d95edf800b4a9de1b16a0ccaffabf1a6 Mon Sep 17 00:00:00 2001 From: Uros Bojanic Date: Fri, 13 Sep 2024 19:04:49 +0200 Subject: [PATCH 027/189] [SPARK-49392][SQL] Catch errors when failing to write to external data source ### What changes were proposed in this pull request? Catch various exceptions thrown by the data source API, when failing to write to a data source, and rethrow `externalDataSourceException` to provide a more friendly error message for the user. ### Why are the changes needed? To catch non-fatal exceptions when failing to save the results of a query into an external data source (for example: `com.crealytics.spark.excel`). ### Does this PR introduce _any_ user-facing change? Yes, error messages when failing to write to an external data source should now be more user-friendly. ### How was this patch tested? New test in `ExplainSuite`. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47873 from uros-db/external-data-source. Authored-by: Uros Bojanic Signed-off-by: Max Gekk --- .../resources/error/error-conditions.json | 6 +++++ .../main/resources/error/error-states.json | 6 +++++ .../sql/errors/QueryCompilationErrors.scala | 8 ++++++ .../SaveIntoDataSourceCommand.scala | 26 ++++++++++++++++--- .../errors/QueryCompilationErrorsSuite.scala | 20 ++++++++++++++ 5 files changed, 63 insertions(+), 3 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 229da4fa17de7..0a9dcd52ea831 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1045,6 +1045,12 @@ ], "sqlState" : "42710" }, + "DATA_SOURCE_EXTERNAL_ERROR" : { + "message" : [ + "Encountered error when saving to external data source." + ], + "sqlState" : "KD00F" + }, "DATA_SOURCE_NOT_EXIST" : { "message" : [ "Data source '' not found. Please make sure the data source is registered." diff --git a/common/utils/src/main/resources/error/error-states.json b/common/utils/src/main/resources/error/error-states.json index c369db3f65058..edba6e1d43216 100644 --- a/common/utils/src/main/resources/error/error-states.json +++ b/common/utils/src/main/resources/error/error-states.json @@ -7417,6 +7417,12 @@ "standard": "N", "usedBy": ["Databricks"] }, + "KD00F": { + "description": "external data source failure", + "origin": "Databricks", + "standard": "N", + "usedBy": ["Databricks"] + }, "P0000": { "description": "procedural logic error", "origin": "PostgreSQL", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index fa8ea2f5289fa..e4c8c76e958f8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -3959,6 +3959,14 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat messageParameters = Map("provider" -> name)) } + def externalDataSourceException(cause: Throwable): Throwable = { + new AnalysisException( + errorClass = "DATA_SOURCE_EXTERNAL_ERROR", + messageParameters = Map(), + cause = Some(cause) + ) + } + def foundMultipleDataSources(provider: String): Throwable = { new AnalysisException( errorClass = "FOUND_MULTIPLE_DATA_SOURCES", diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SaveIntoDataSourceCommand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SaveIntoDataSourceCommand.scala index 5423232db4293..e44f1d35e9cdf 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SaveIntoDataSourceCommand.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SaveIntoDataSourceCommand.scala @@ -19,12 +19,14 @@ package org.apache.spark.sql.execution.datasources import scala.util.control.NonFatal +import org.apache.spark.SparkThrowable import org.apache.spark.sql.{Dataset, Row, SaveMode, SparkSession} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{CTEInChildren, CTERelationDef, LogicalPlan, WithCTE} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes +import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.LeafRunnableCommand -import org.apache.spark.sql.sources.CreatableRelationProvider +import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider} /** * Saves the results of `query` in to a data source. @@ -44,8 +46,26 @@ case class SaveIntoDataSourceCommand( override def innerChildren: Seq[QueryPlan[_]] = Seq(query) override def run(sparkSession: SparkSession): Seq[Row] = { - val relation = dataSource.createRelation( - sparkSession.sqlContext, mode, options, Dataset.ofRows(sparkSession, query)) + var relation: BaseRelation = null + + try { + relation = dataSource.createRelation( + sparkSession.sqlContext, mode, options, Dataset.ofRows(sparkSession, query)) + } catch { + case e: SparkThrowable => + // We should avoid wrapping `SparkThrowable` exceptions into another `AnalysisException`. + throw e + case e @ (_: NullPointerException | _: MatchError | _: ArrayIndexOutOfBoundsException) => + // These are some of the exceptions thrown by the data source API. We catch these + // exceptions here and rethrow QueryCompilationErrors.externalDataSourceException to + // provide a more friendly error message for the user. This list is not exhaustive. + throw QueryCompilationErrors.externalDataSourceException(e) + case e: Throwable => + // For other exceptions, just rethrow it, since we don't have enough information to + // provide a better error message for the user at the moment. We may want to further + // improve the error message handling in the future. + throw e + } try { val logicalRelation = LogicalRelation(relation, toAttributes(relation.schema), None, false) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala index 47a6143bad1d7..370c118de9a93 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala @@ -24,6 +24,7 @@ import org.apache.spark.sql._ import org.apache.spark.sql.api.java.{UDF1, UDF2, UDF23Test} import org.apache.spark.sql.catalyst.expressions.{Coalesce, Literal, UnsafeRow} import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.execution.datasources.SaveIntoDataSourceCommand import org.apache.spark.sql.execution.datasources.parquet.SparkToParquetSchemaConverter import org.apache.spark.sql.expressions.SparkUserDefinedFunction import org.apache.spark.sql.functions._ @@ -926,6 +927,25 @@ class QueryCompilationErrorsSuite }) } } + + test("Catch and log errors when failing to write to external data source") { + val password = "MyPassWord" + val token = "MyToken" + val value = "value" + val options = Map("password" -> password, "token" -> token, "key" -> value) + val query = spark.range(10).logicalPlan + val cmd = SaveIntoDataSourceCommand(query, null, options, SaveMode.Overwrite) + + checkError( + exception = intercept[AnalysisException] { + cmd.run(spark) + }, + condition = "DATA_SOURCE_EXTERNAL_ERROR", + sqlState = "KD00F", + parameters = Map.empty + ) + } + } class MyCastToString extends SparkUserDefinedFunction( From 08a26bb56cfb48f27c68a79be1e15bc4c9e466e0 Mon Sep 17 00:00:00 2001 From: viktorluc-db Date: Fri, 13 Sep 2024 19:16:28 +0200 Subject: [PATCH 028/189] [SPARK-48779][SQL][TESTS] Improve collation support testing - add golden files ### What changes were proposed in this pull request? Moving certain tests to collations.sql and generating golden files, as said in the PR https://github.com/apache/spark/pull/47620. ### Why are the changes needed? For testing. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? These are tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47828 from viktorluc-db/GoldenFiles. Authored-by: viktorluc-db Signed-off-by: Max Gekk --- .../analyzer-results/collations.sql.out | 1831 ++++++++- .../resources/sql-tests/inputs/collations.sql | 322 +- .../sql-tests/results/collations.sql.out | 3537 ++++++++++++++++- 3 files changed, 5615 insertions(+), 75 deletions(-) diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out index 14ac67eb93a32..83c9ebfef4b25 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out @@ -444,21 +444,1493 @@ DropTable false, false -- !query -create table t5(str string collate utf8_binary, delimiter string collate utf8_lcase, partNum int) using parquet +create table t5(s string, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet -- !query analysis CreateDataSourceTableCommand `spark_catalog`.`default`.`t5`, false -- !query -insert into t5 values('11AB12AB13', 'AB', 2) +insert into t5 values ('Spark', 'Spark', 'SQL') -- !query analysis -InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [str, delimiter, partNum] -+- Project [cast(col1#x as string) AS str#x, cast(col2#x as string collate UTF8_LCASE) AS delimiter#x, cast(col3#x as int) AS partNum#x] +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] +- LocalRelation [col1#x, col2#x, col3#x] -- !query -select split_part(str, delimiter, partNum) from t5 +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaAAaA') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaA') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaAaaAaaAaAaaAaaAaA') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('bbAbaAbA', 'bbAbAAbA', 'a') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('İo', 'İo', 'İo') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('İo', 'İo', 'i̇o') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('efd2', 'efd2', 'efd2') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('Hello, world! Nice day.', 'Hello, world! Nice day.', 'Hello, world! Nice day.') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('Something else. Nothing here.', 'Something else. Nothing here.', 'Something else. Nothing here.') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('kitten', 'kitten', 'sitTing') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('abc', 'abc', 'abc') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +insert into t5 values ('abcdcba', 'abcdcba', 'aBcDCbA') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t5, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t5], Append, `spark_catalog`.`default`.`t5`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t5), [s, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS s#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +create table t6(ascii long) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t6`, false + + +-- !query +insert into t6 values (97) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t6, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t6], Append, `spark_catalog`.`default`.`t6`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t6), [ascii] ++- Project [cast(col1#x as bigint) AS ascii#xL] + +- LocalRelation [col1#x] + + +-- !query +insert into t6 values (66) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t6, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t6], Append, `spark_catalog`.`default`.`t6`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t6), [ascii] ++- Project [cast(col1#x as bigint) AS ascii#xL] + +- LocalRelation [col1#x] + + +-- !query +create table t7(ascii double) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t7`, false + + +-- !query +insert into t7 values (97.52143) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t7, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t7], Append, `spark_catalog`.`default`.`t7`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t7), [ascii] ++- Project [cast(col1#x as double) AS ascii#x] + +- LocalRelation [col1#x] + + +-- !query +insert into t7 values (66.421) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t7, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t7], Append, `spark_catalog`.`default`.`t7`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t7), [ascii] ++- Project [cast(col1#x as double) AS ascii#x] + +- LocalRelation [col1#x] + + +-- !query +create table t8(format string collate utf8_binary, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t8`, false + + +-- !query +insert into t8 values ('%s%s', 'abCdE', 'abCdE') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t8, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t8], Append, `spark_catalog`.`default`.`t8`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t8), [format, utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS format#x, cast(col2#x as string) AS utf8_binary#x, cast(col3#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +create table t9(num long) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t9`, false + + +-- !query +insert into t9 values (97) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t9, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t9], Append, `spark_catalog`.`default`.`t9`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t9), [num] ++- Project [cast(col1#x as bigint) AS num#xL] + +- LocalRelation [col1#x] + + +-- !query +insert into t9 values (66) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t9, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t9], Append, `spark_catalog`.`default`.`t9`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t9), [num] ++- Project [cast(col1#x as bigint) AS num#xL] + +- LocalRelation [col1#x] + + +-- !query +create table t10(utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t10`, false + + +-- !query +insert into t10 values ('aaAaAAaA', 'aaAaaAaA') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t10, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t10], Append, `spark_catalog`.`default`.`t10`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t10), [utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS utf8_binary#x, cast(col2#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +insert into t10 values ('efd2', 'efd2') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t10, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t10], Append, `spark_catalog`.`default`.`t10`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t10), [utf8_binary, utf8_lcase] ++- Project [cast(col1#x as string) AS utf8_binary#x, cast(col2#x as string collate UTF8_LCASE) AS utf8_lcase#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +select concat_ws(' ', utf8_lcase, utf8_lcase) from t5 +-- !query analysis +Project [concat_ws(cast( as string collate UTF8_LCASE), utf8_lcase#x, utf8_lcase#x) AS concat_ws( , utf8_lcase, utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select concat_ws(' ', utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select concat_ws(' ' collate utf8_binary, utf8_binary, 'SQL' collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select concat_ws(' ' collate utf8_lcase, utf8_binary, 'SQL' collate utf8_lcase) from t5 +-- !query analysis +Project [concat_ws(collate( , utf8_lcase), cast(utf8_binary#x as string collate UTF8_LCASE), collate(SQL, utf8_lcase)) AS concat_ws(collate( , utf8_lcase), utf8_binary, collate(SQL, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select concat_ws(',', utf8_lcase, 'word'), concat_ws(',', utf8_binary, 'word') from t5 +-- !query analysis +Project [concat_ws(cast(, as string collate UTF8_LCASE), utf8_lcase#x, cast(word as string collate UTF8_LCASE)) AS concat_ws(,, utf8_lcase, word)#x, concat_ws(,, utf8_binary#x, word) AS concat_ws(,, utf8_binary, word)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select concat_ws(',', utf8_lcase, 'word' collate utf8_binary), concat_ws(',', utf8_binary, 'word' collate utf8_lcase) from t5 +-- !query analysis +Project [concat_ws(,, cast(utf8_lcase#x as string), collate(word, utf8_binary)) AS concat_ws(,, utf8_lcase, collate(word, utf8_binary))#x, concat_ws(cast(, as string collate UTF8_LCASE), cast(utf8_binary#x as string collate UTF8_LCASE), collate(word, utf8_lcase)) AS concat_ws(,, utf8_binary, collate(word, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select elt(2, s, utf8_binary) from t5 +-- !query analysis +Project [elt(2, s#x, utf8_binary#x, false) AS elt(2, s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select elt(2, utf8_binary, utf8_lcase, s) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [elt(1, collate(utf8_binary#x, utf8_binary), collate(utf8_lcase#x, utf8_binary), false) AS elt(1, collate(utf8_binary, utf8_binary), collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase) from t5 +-- !query analysis +Project [elt(1, collate(utf8_binary#x, utf8_binary), cast(utf8_lcase#x as string), false) AS elt(1, collate(utf8_binary, utf8_binary), utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select elt(1, utf8_binary, 'word'), elt(1, utf8_lcase, 'word') from t5 +-- !query analysis +Project [elt(1, utf8_binary#x, word, false) AS elt(1, utf8_binary, word)#x, elt(1, utf8_lcase#x, cast(word as string collate UTF8_LCASE), false) AS elt(1, utf8_lcase, word)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select elt(1, utf8_binary, 'word' collate utf8_lcase), elt(1, utf8_lcase, 'word' collate utf8_binary) from t5 +-- !query analysis +Project [elt(1, cast(utf8_binary#x as string collate UTF8_LCASE), collate(word, utf8_lcase), false) AS elt(1, utf8_binary, collate(word, utf8_lcase))#x, elt(1, cast(utf8_lcase#x as string), collate(word, utf8_binary), false) AS elt(1, utf8_lcase, collate(word, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select split_part(utf8_binary, utf8_lcase, 3) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select split_part(s, utf8_binary, 1) from t5 +-- !query analysis +Project [split_part(s#x, utf8_binary#x, 1) AS split_part(s, utf8_binary, 1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select split_part(utf8_binary collate utf8_binary, s collate utf8_lcase, 1) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select split_part(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [split_part(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), 2) AS split_part(utf8_binary, collate(utf8_lcase, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select split_part(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query analysis +Project [split_part(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), 2) AS split_part(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5 +-- !query analysis +Project [split_part(utf8_binary#x, a, 3) AS split_part(utf8_binary, a, 3)#x, split_part(utf8_lcase#x, cast(a as string collate UTF8_LCASE), 3) AS split_part(utf8_lcase, a, 3)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select split_part(utf8_binary, 'a' collate utf8_lcase, 3), split_part(utf8_lcase, 'a' collate utf8_binary, 3) from t5 +-- !query analysis +Project [split_part(cast(utf8_binary#x as string collate UTF8_LCASE), collate(a, utf8_lcase), 3) AS split_part(utf8_binary, collate(a, utf8_lcase), 3)#x, split_part(cast(utf8_lcase#x as string), collate(a, utf8_binary), 3) AS split_part(utf8_lcase, collate(a, utf8_binary), 3)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select contains(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select contains(s, utf8_binary) from t5 +-- !query analysis +Project [Contains(s#x, utf8_binary#x) AS contains(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select contains(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select contains(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [Contains(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS contains(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select contains(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [Contains(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS contains(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5 +-- !query analysis +Project [Contains(utf8_binary#x, a) AS contains(utf8_binary, a)#x, Contains(utf8_lcase#x, cast(a as string collate UTF8_LCASE)) AS contains(utf8_lcase, a)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select contains(utf8_binary, 'AaAA' collate utf8_lcase), contains(utf8_lcase, 'AAa' collate utf8_binary) from t5 +-- !query analysis +Project [Contains(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase)) AS contains(utf8_binary, collate(AaAA, utf8_lcase))#x, Contains(cast(utf8_lcase#x as string), collate(AAa, utf8_binary)) AS contains(utf8_lcase, collate(AAa, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substring_index(utf8_binary, utf8_lcase, 2) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select substring_index(s, utf8_binary,1) from t5 +-- !query analysis +Project [substring_index(s#x, utf8_binary#x, 1) AS substring_index(s, utf8_binary, 1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substring_index(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select substring_index(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [substring_index(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), 2) AS substring_index(utf8_binary, collate(utf8_lcase, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substring_index(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query analysis +Project [substring_index(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), 2) AS substring_index(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5 +-- !query analysis +Project [substring_index(utf8_binary#x, a, 2) AS substring_index(utf8_binary, a, 2)#x, substring_index(utf8_lcase#x, cast(a as string collate UTF8_LCASE), 2) AS substring_index(utf8_lcase, a, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substring_index(utf8_binary, 'AaAA' collate utf8_lcase, 2), substring_index(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5 +-- !query analysis +Project [substring_index(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase), 2) AS substring_index(utf8_binary, collate(AaAA, utf8_lcase), 2)#x, substring_index(cast(utf8_lcase#x as string), collate(AAa, utf8_binary), 2) AS substring_index(utf8_lcase, collate(AAa, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select instr(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select instr(s, utf8_binary) from t5 +-- !query analysis +Project [instr(s#x, utf8_binary#x) AS instr(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select instr(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select instr(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [instr(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS instr(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select instr(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [instr(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS instr(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5 +-- !query analysis +Project [instr(utf8_binary#x, a) AS instr(utf8_binary, a)#x, instr(utf8_lcase#x, cast(a as string collate UTF8_LCASE)) AS instr(utf8_lcase, a)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select instr(utf8_binary, 'AaAA' collate utf8_lcase), instr(utf8_lcase, 'AAa' collate utf8_binary) from t5 +-- !query analysis +Project [instr(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase)) AS instr(utf8_binary, collate(AaAA, utf8_lcase))#x, instr(cast(utf8_lcase#x as string), collate(AAa, utf8_binary)) AS instr(utf8_lcase, collate(AAa, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select find_in_set(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select find_in_set(s, utf8_binary) from t5 +-- !query analysis +Project [find_in_set(s#x, utf8_binary#x) AS find_in_set(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select find_in_set(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select find_in_set(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [find_in_set(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS find_in_set(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select find_in_set(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [find_in_set(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS find_in_set(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o'), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o') from t5 +-- !query analysis +Project [find_in_set(utf8_binary#x, aaAaaAaA,i̇o) AS find_in_set(utf8_binary, aaAaaAaA,i̇o)#x, find_in_set(utf8_lcase#x, cast(aaAaaAaA,i̇o as string collate UTF8_LCASE)) AS find_in_set(utf8_lcase, aaAaaAaA,i̇o)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o' collate utf8_lcase), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o' collate utf8_binary) from t5 +-- !query analysis +Project [find_in_set(cast(utf8_binary#x as string collate UTF8_LCASE), collate(aaAaaAaA,i̇o, utf8_lcase)) AS find_in_set(utf8_binary, collate(aaAaaAaA,i̇o, utf8_lcase))#x, find_in_set(cast(utf8_lcase#x as string), collate(aaAaaAaA,i̇o, utf8_binary)) AS find_in_set(utf8_lcase, collate(aaAaaAaA,i̇o, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select startswith(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select startswith(s, utf8_binary) from t5 +-- !query analysis +Project [StartsWith(s#x, utf8_binary#x) AS startswith(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select startswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select startswith(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [StartsWith(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS startswith(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select startswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [StartsWith(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS startswith(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5 +-- !query analysis +Project [StartsWith(utf8_binary#x, aaAaaAaA) AS startswith(utf8_binary, aaAaaAaA)#x, StartsWith(utf8_lcase#x, cast(aaAaaAaA as string collate UTF8_LCASE)) AS startswith(utf8_lcase, aaAaaAaA)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select startswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), startswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5 +-- !query analysis +Project [StartsWith(cast(utf8_binary#x as string collate UTF8_LCASE), collate(aaAaaAaA, utf8_lcase)) AS startswith(utf8_binary, collate(aaAaaAaA, utf8_lcase))#x, StartsWith(cast(utf8_lcase#x as string), collate(aaAaaAaA, utf8_binary)) AS startswith(utf8_lcase, collate(aaAaaAaA, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select translate(utf8_lcase, utf8_lcase, '12345') from t5 +-- !query analysis +Project [translate(utf8_lcase#x, utf8_lcase#x, cast(12345 as string collate UTF8_LCASE)) AS translate(utf8_lcase, utf8_lcase, 12345)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select translate(utf8_binary, utf8_lcase, '12345') from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select translate(utf8_binary, 'aBc' collate utf8_lcase, '12345' collate utf8_binary) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" + } +} + + +-- !query +select translate(utf8_binary, 'SQL' collate utf8_lcase, '12345' collate utf8_lcase) from t5 +-- !query analysis +Project [translate(cast(utf8_binary#x as string collate UTF8_LCASE), collate(SQL, utf8_lcase), collate(12345, utf8_lcase)) AS translate(utf8_binary, collate(SQL, utf8_lcase), collate(12345, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5 +-- !query analysis +Project [translate(utf8_lcase#x, cast(aaAaaAaA as string collate UTF8_LCASE), cast(12345 as string collate UTF8_LCASE)) AS translate(utf8_lcase, aaAaaAaA, 12345)#x, translate(utf8_binary#x, aaAaaAaA, 12345) AS translate(utf8_binary, aaAaaAaA, 12345)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select translate(utf8_lcase, 'aBc' collate utf8_binary, '12345'), translate(utf8_binary, 'aBc' collate utf8_lcase, '12345') from t5 +-- !query analysis +Project [translate(cast(utf8_lcase#x as string), collate(aBc, utf8_binary), 12345) AS translate(utf8_lcase, collate(aBc, utf8_binary), 12345)#x, translate(cast(utf8_binary#x as string collate UTF8_LCASE), collate(aBc, utf8_lcase), cast(12345 as string collate UTF8_LCASE)) AS translate(utf8_binary, collate(aBc, utf8_lcase), 12345)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select replace(utf8_binary, utf8_lcase, 'abc') from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select replace(s, utf8_binary, 'abc') from t5 +-- !query analysis +Project [replace(s#x, utf8_binary#x, abc) AS replace(s, utf8_binary, abc)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select replace(utf8_binary collate utf8_binary, s collate utf8_lcase, 'abc') from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select replace(utf8_binary, utf8_lcase collate utf8_binary, 'abc') from t5 +-- !query analysis +Project [replace(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), abc) AS replace(utf8_binary, collate(utf8_lcase, utf8_binary), abc)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select replace(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 'abc') from t5 +-- !query analysis +Project [replace(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), cast(abc as string collate UTF8_LCASE)) AS replace(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase), abc)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5 +-- !query analysis +Project [replace(utf8_binary#x, aaAaaAaA, abc) AS replace(utf8_binary, aaAaaAaA, abc)#x, replace(utf8_lcase#x, cast(aaAaaAaA as string collate UTF8_LCASE), cast(abc as string collate UTF8_LCASE)) AS replace(utf8_lcase, aaAaaAaA, abc)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select replace(utf8_binary, 'aaAaaAaA' collate utf8_lcase, 'abc'), replace(utf8_lcase, 'aaAaaAaA' collate utf8_binary, 'abc') from t5 +-- !query analysis +Project [replace(cast(utf8_binary#x as string collate UTF8_LCASE), collate(aaAaaAaA, utf8_lcase), cast(abc as string collate UTF8_LCASE)) AS replace(utf8_binary, collate(aaAaaAaA, utf8_lcase), abc)#x, replace(cast(utf8_lcase#x as string), collate(aaAaaAaA, utf8_binary), abc) AS replace(utf8_lcase, collate(aaAaaAaA, utf8_binary), abc)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select endswith(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select endswith(s, utf8_binary) from t5 +-- !query analysis +Project [EndsWith(s#x, utf8_binary#x) AS endswith(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select endswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select endswith(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [EndsWith(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS endswith(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select endswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [EndsWith(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS endswith(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5 +-- !query analysis +Project [EndsWith(utf8_binary#x, aaAaaAaA) AS endswith(utf8_binary, aaAaaAaA)#x, EndsWith(utf8_lcase#x, cast(aaAaaAaA as string collate UTF8_LCASE)) AS endswith(utf8_lcase, aaAaaAaA)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select endswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), endswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5 +-- !query analysis +Project [EndsWith(cast(utf8_binary#x as string collate UTF8_LCASE), collate(aaAaaAaA, utf8_lcase)) AS endswith(utf8_binary, collate(aaAaaAaA, utf8_lcase))#x, EndsWith(cast(utf8_lcase#x as string), collate(aaAaaAaA, utf8_binary)) AS endswith(utf8_lcase, collate(aaAaaAaA, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select repeat(utf8_binary, 3), repeat(utf8_lcase, 2) from t5 +-- !query analysis +Project [repeat(utf8_binary#x, 3) AS repeat(utf8_binary, 3)#x, repeat(utf8_lcase#x, 2) AS repeat(utf8_lcase, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select repeat(utf8_binary collate utf8_lcase, 3), repeat(utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [repeat(collate(utf8_binary#x, utf8_lcase), 3) AS repeat(collate(utf8_binary, utf8_lcase), 3)#x, repeat(collate(utf8_lcase#x, utf8_binary), 2) AS repeat(collate(utf8_lcase, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select ascii(utf8_binary), ascii(utf8_lcase) from t5 +-- !query analysis +Project [ascii(utf8_binary#x) AS ascii(utf8_binary)#x, ascii(utf8_lcase#x) AS ascii(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select ascii(utf8_binary collate utf8_lcase), ascii(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [ascii(collate(utf8_binary#x, utf8_lcase)) AS ascii(collate(utf8_binary, utf8_lcase))#x, ascii(collate(utf8_lcase#x, utf8_binary)) AS ascii(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select unbase64(utf8_binary), unbase64(utf8_lcase) from t10 +-- !query analysis +Project [unbase64(utf8_binary#x, false) AS unbase64(utf8_binary)#x, unbase64(utf8_lcase#x, false) AS unbase64(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t10 + +- Relation spark_catalog.default.t10[utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select unbase64(utf8_binary collate utf8_lcase), unbase64(utf8_lcase collate utf8_binary) from t10 +-- !query analysis +Project [unbase64(collate(utf8_binary#x, utf8_lcase), false) AS unbase64(collate(utf8_binary, utf8_lcase))#x, unbase64(collate(utf8_lcase#x, utf8_binary), false) AS unbase64(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t10 + +- Relation spark_catalog.default.t10[utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select chr(ascii) from t6 +-- !query analysis +Project [chr(ascii#xL) AS chr(ascii)#x] ++- SubqueryAlias spark_catalog.default.t6 + +- Relation spark_catalog.default.t6[ascii#xL] parquet + + +-- !query +select base64(utf8_binary), base64(utf8_lcase) from t5 +-- !query analysis +Project [base64(cast(utf8_binary#x as binary)) AS base64(utf8_binary)#x, base64(cast(utf8_lcase#x as binary)) AS base64(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select base64(utf8_binary collate utf8_lcase), base64(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [base64(cast(collate(utf8_binary#x, utf8_lcase) as binary)) AS base64(collate(utf8_binary, utf8_lcase))#x, base64(cast(collate(utf8_lcase#x, utf8_binary) as binary)) AS base64(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select decode(encode(utf8_binary, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase, 'utf-8'), 'utf-8') from t5 +-- !query analysis +Project [decode(encode(utf8_binary#x, utf-8), utf-8) AS decode(encode(utf8_binary, utf-8), utf-8)#x, decode(encode(utf8_lcase#x, utf-8), utf-8) AS decode(encode(utf8_lcase, utf-8), utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select decode(encode(utf8_binary collate utf8_lcase, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase collate utf8_binary, 'utf-8'), 'utf-8') from t5 +-- !query analysis +Project [decode(encode(collate(utf8_binary#x, utf8_lcase), utf-8), utf-8) AS decode(encode(collate(utf8_binary, utf8_lcase), utf-8), utf-8)#x, decode(encode(collate(utf8_lcase#x, utf8_binary), utf-8), utf-8) AS decode(encode(collate(utf8_lcase, utf8_binary), utf-8), utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select format_number(ascii, '###.###') from t7 +-- !query analysis +Project [format_number(ascii#x, ###.###) AS format_number(ascii, ###.###)#x] ++- SubqueryAlias spark_catalog.default.t7 + +- Relation spark_catalog.default.t7[ascii#x] parquet + + +-- !query +select format_number(ascii, '###.###' collate utf8_lcase) from t7 +-- !query analysis +Project [format_number(ascii#x, collate(###.###, utf8_lcase)) AS format_number(ascii, collate(###.###, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t7 + +- Relation spark_catalog.default.t7[ascii#x] parquet + + +-- !query +select encode(utf8_binary, 'utf-8'), encode(utf8_lcase, 'utf-8') from t5 +-- !query analysis +Project [encode(utf8_binary#x, utf-8) AS encode(utf8_binary, utf-8)#x, encode(utf8_lcase#x, utf-8) AS encode(utf8_lcase, utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select encode(utf8_binary collate utf8_lcase, 'utf-8'), encode(utf8_lcase collate utf8_binary, 'utf-8') from t5 +-- !query analysis +Project [encode(collate(utf8_binary#x, utf8_lcase), utf-8) AS encode(collate(utf8_binary, utf8_lcase), utf-8)#x, encode(collate(utf8_lcase#x, utf8_binary), utf-8) AS encode(collate(utf8_lcase, utf8_binary), utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select to_binary(utf8_binary, 'utf-8'), to_binary(utf8_lcase, 'utf-8') from t5 +-- !query analysis +Project [to_binary(utf8_binary#x, Some(utf-8), false) AS to_binary(utf8_binary, utf-8)#x, to_binary(utf8_lcase#x, Some(utf-8), false) AS to_binary(utf8_lcase, utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select to_binary(utf8_binary collate utf8_lcase, 'utf-8'), to_binary(utf8_lcase collate utf8_binary, 'utf-8') from t5 +-- !query analysis +Project [to_binary(collate(utf8_binary#x, utf8_lcase), Some(utf-8), false) AS to_binary(collate(utf8_binary, utf8_lcase), utf-8)#x, to_binary(collate(utf8_lcase#x, utf8_binary), Some(utf-8), false) AS to_binary(collate(utf8_lcase, utf8_binary), utf-8)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select sentences(utf8_binary), sentences(utf8_lcase) from t5 +-- !query analysis +Project [sentences(utf8_binary#x, , ) AS sentences(utf8_binary, , )#x, sentences(utf8_lcase#x, , ) AS sentences(utf8_lcase, , )#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select sentences(utf8_binary collate utf8_lcase), sentences(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [sentences(collate(utf8_binary#x, utf8_lcase), , ) AS sentences(collate(utf8_binary, utf8_lcase), , )#x, sentences(collate(utf8_lcase#x, utf8_binary), , ) AS sentences(collate(utf8_lcase, utf8_binary), , )#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select upper(utf8_binary), upper(utf8_lcase) from t5 +-- !query analysis +Project [upper(utf8_binary#x) AS upper(utf8_binary)#x, upper(utf8_lcase#x) AS upper(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select upper(utf8_binary collate utf8_lcase), upper(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [upper(collate(utf8_binary#x, utf8_lcase)) AS upper(collate(utf8_binary, utf8_lcase))#x, upper(collate(utf8_lcase#x, utf8_binary)) AS upper(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select lower(utf8_binary), lower(utf8_lcase) from t5 +-- !query analysis +Project [lower(utf8_binary#x) AS lower(utf8_binary)#x, lower(utf8_lcase#x) AS lower(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select lower(utf8_binary collate utf8_lcase), lower(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [lower(collate(utf8_binary#x, utf8_lcase)) AS lower(collate(utf8_binary, utf8_lcase))#x, lower(collate(utf8_lcase#x, utf8_binary)) AS lower(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select initcap(utf8_binary), initcap(utf8_lcase) from t5 +-- !query analysis +Project [initcap(utf8_binary#x) AS initcap(utf8_binary)#x, initcap(utf8_lcase#x) AS initcap(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select initcap(utf8_binary collate utf8_lcase), initcap(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [initcap(collate(utf8_binary#x, utf8_lcase)) AS initcap(collate(utf8_binary, utf8_lcase))#x, initcap(collate(utf8_lcase#x, utf8_binary)) AS initcap(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select overlay(utf8_binary, utf8_lcase, 2) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select overlay(s, utf8_binary,1) from t5 +-- !query analysis +Project [overlay(s#x, utf8_binary#x, 1, -1) AS overlay(s, utf8_binary, 1, -1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select overlay(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select overlay(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [overlay(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), 2, -1) AS overlay(utf8_binary, collate(utf8_lcase, utf8_binary), 2, -1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select overlay(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query analysis +Project [overlay(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), 2, -1) AS overlay(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase), 2, -1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select overlay(utf8_binary, 'a', 2), overlay(utf8_lcase, 'a', 2) from t5 +-- !query analysis +Project [overlay(utf8_binary#x, a, 2, -1) AS overlay(utf8_binary, a, 2, -1)#x, overlay(utf8_lcase#x, cast(a as string collate UTF8_LCASE), 2, -1) AS overlay(utf8_lcase, a, 2, -1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select overlay(utf8_binary, 'AaAA' collate utf8_lcase, 2), overlay(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5 +-- !query analysis +Project [overlay(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase), 2, -1) AS overlay(utf8_binary, collate(AaAA, utf8_lcase), 2, -1)#x, overlay(cast(utf8_lcase#x as string), collate(AAa, utf8_binary), 2, -1) AS overlay(utf8_lcase, collate(AAa, utf8_binary), 2, -1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select format_string(format, utf8_binary, utf8_lcase) from t8 +-- !query analysis +Project [format_string(format#x, utf8_binary#x, utf8_lcase#x) AS format_string(format, utf8_binary, utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t8 + +- Relation spark_catalog.default.t8[format#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select format_string(format collate utf8_lcase, utf8_lcase, utf8_binary collate utf8_lcase, 3), format_string(format, utf8_lcase collate utf8_binary, utf8_binary) from t8 +-- !query analysis +Project [format_string(collate(format#x, utf8_lcase), utf8_lcase#x, collate(utf8_binary#x, utf8_lcase), 3) AS format_string(collate(format, utf8_lcase), utf8_lcase, collate(utf8_binary, utf8_lcase), 3)#x, format_string(format#x, collate(utf8_lcase#x, utf8_binary), utf8_binary#x) AS format_string(format, collate(utf8_lcase, utf8_binary), utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t8 + +- Relation spark_catalog.default.t8[format#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select format_string(format, utf8_binary, utf8_lcase) from t8 +-- !query analysis +Project [format_string(format#x, utf8_binary#x, utf8_lcase#x) AS format_string(format, utf8_binary, utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t8 + +- Relation spark_catalog.default.t8[format#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select soundex(utf8_binary), soundex(utf8_lcase) from t5 +-- !query analysis +Project [soundex(utf8_binary#x) AS soundex(utf8_binary)#x, soundex(utf8_lcase#x) AS soundex(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select soundex(utf8_binary collate utf8_lcase), soundex(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [soundex(collate(utf8_binary#x, utf8_lcase)) AS soundex(collate(utf8_binary, utf8_lcase))#x, soundex(collate(utf8_lcase#x, utf8_binary)) AS soundex(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select length(utf8_binary), length(utf8_lcase) from t5 +-- !query analysis +Project [length(utf8_binary#x) AS length(utf8_binary)#x, length(utf8_lcase#x) AS length(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select length(utf8_binary collate utf8_lcase), length(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [length(collate(utf8_binary#x, utf8_lcase)) AS length(collate(utf8_binary, utf8_lcase))#x, length(collate(utf8_lcase#x, utf8_binary)) AS length(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select bit_length(utf8_binary), bit_length(utf8_lcase) from t5 +-- !query analysis +Project [bit_length(utf8_binary#x) AS bit_length(utf8_binary)#x, bit_length(utf8_lcase#x) AS bit_length(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select bit_length(utf8_binary collate utf8_lcase), bit_length(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [bit_length(collate(utf8_binary#x, utf8_lcase)) AS bit_length(collate(utf8_binary, utf8_lcase))#x, bit_length(collate(utf8_lcase#x, utf8_binary)) AS bit_length(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select octet_length(utf8_binary), octet_length(utf8_lcase) from t5 +-- !query analysis +Project [octet_length(utf8_binary#x) AS octet_length(utf8_binary)#x, octet_length(utf8_lcase#x) AS octet_length(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select octet_length(utf8_binary collate utf8_lcase), octet_length(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [octet_length(collate(utf8_binary#x, utf8_lcase)) AS octet_length(collate(utf8_binary, utf8_lcase))#x, octet_length(collate(utf8_lcase#x, utf8_binary)) AS octet_length(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select luhn_check(num) from t9 +-- !query analysis +Project [luhn_check(cast(num#xL as string)) AS luhn_check(num)#x] ++- SubqueryAlias spark_catalog.default.t9 + +- Relation spark_catalog.default.t9[num#xL] parquet + + +-- !query +select levenshtein(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select levenshtein(s, utf8_binary) from t5 +-- !query analysis +Project [levenshtein(s#x, utf8_binary#x, None) AS levenshtein(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select levenshtein(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select levenshtein(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [levenshtein(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), None) AS levenshtein(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select levenshtein(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [levenshtein(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), None) AS levenshtein(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select levenshtein(utf8_binary, 'a'), levenshtein(utf8_lcase, 'a') from t5 +-- !query analysis +Project [levenshtein(utf8_binary#x, a, None) AS levenshtein(utf8_binary, a)#x, levenshtein(utf8_lcase#x, cast(a as string collate UTF8_LCASE), None) AS levenshtein(utf8_lcase, a)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select levenshtein(utf8_binary, 'AaAA' collate utf8_lcase, 3), levenshtein(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5 +-- !query analysis +Project [levenshtein(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase), Some(3)) AS levenshtein(utf8_binary, collate(AaAA, utf8_lcase), 3)#x, levenshtein(cast(utf8_lcase#x as string), collate(AAa, utf8_binary), Some(4)) AS levenshtein(utf8_lcase, collate(AAa, utf8_binary), 4)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select is_valid_utf8(utf8_binary), is_valid_utf8(utf8_lcase) from t5 +-- !query analysis +Project [is_valid_utf8(utf8_binary#x) AS is_valid_utf8(utf8_binary)#x, is_valid_utf8(utf8_lcase#x) AS is_valid_utf8(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select is_valid_utf8(utf8_binary collate utf8_lcase), is_valid_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [is_valid_utf8(collate(utf8_binary#x, utf8_lcase)) AS is_valid_utf8(collate(utf8_binary, utf8_lcase))#x, is_valid_utf8(collate(utf8_lcase#x, utf8_binary)) AS is_valid_utf8(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select make_valid_utf8(utf8_binary), make_valid_utf8(utf8_lcase) from t5 +-- !query analysis +Project [make_valid_utf8(utf8_binary#x) AS make_valid_utf8(utf8_binary)#x, make_valid_utf8(utf8_lcase#x) AS make_valid_utf8(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select make_valid_utf8(utf8_binary collate utf8_lcase), make_valid_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [make_valid_utf8(collate(utf8_binary#x, utf8_lcase)) AS make_valid_utf8(collate(utf8_binary, utf8_lcase))#x, make_valid_utf8(collate(utf8_lcase#x, utf8_binary)) AS make_valid_utf8(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select validate_utf8(utf8_binary), validate_utf8(utf8_lcase) from t5 +-- !query analysis +Project [validate_utf8(utf8_binary#x) AS validate_utf8(utf8_binary)#x, validate_utf8(utf8_lcase#x) AS validate_utf8(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select validate_utf8(utf8_binary collate utf8_lcase), validate_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [validate_utf8(collate(utf8_binary#x, utf8_lcase)) AS validate_utf8(collate(utf8_binary, utf8_lcase))#x, validate_utf8(collate(utf8_lcase#x, utf8_binary)) AS validate_utf8(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select try_validate_utf8(utf8_binary), try_validate_utf8(utf8_lcase) from t5 +-- !query analysis +Project [try_validate_utf8(utf8_binary#x) AS try_validate_utf8(utf8_binary)#x, try_validate_utf8(utf8_lcase#x) AS try_validate_utf8(utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select try_validate_utf8(utf8_binary collate utf8_lcase), try_validate_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [try_validate_utf8(collate(utf8_binary#x, utf8_lcase)) AS try_validate_utf8(collate(utf8_binary, utf8_lcase))#x, try_validate_utf8(collate(utf8_lcase#x, utf8_binary)) AS try_validate_utf8(collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substr(utf8_binary, 2, 2), substr(utf8_lcase, 2, 2) from t5 +-- !query analysis +Project [substr(utf8_binary#x, 2, 2) AS substr(utf8_binary, 2, 2)#x, substr(utf8_lcase#x, 2, 2) AS substr(utf8_lcase, 2, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select substr(utf8_binary collate utf8_lcase, 2, 2), substr(utf8_lcase collate utf8_binary, 2, 2) from t5 +-- !query analysis +Project [substr(collate(utf8_binary#x, utf8_lcase), 2, 2) AS substr(collate(utf8_binary, utf8_lcase), 2, 2)#x, substr(collate(utf8_lcase#x, utf8_binary), 2, 2) AS substr(collate(utf8_lcase, utf8_binary), 2, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select right(utf8_binary, 2), right(utf8_lcase, 2) from t5 +-- !query analysis +Project [right(utf8_binary#x, 2) AS right(utf8_binary, 2)#x, right(utf8_lcase#x, 2) AS right(utf8_lcase, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select right(utf8_binary collate utf8_lcase, 2), right(utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [right(collate(utf8_binary#x, utf8_lcase), 2) AS right(collate(utf8_binary, utf8_lcase), 2)#x, right(collate(utf8_lcase#x, utf8_binary), 2) AS right(collate(utf8_lcase, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select left(utf8_binary, '2' collate utf8_lcase), left(utf8_lcase, 2) from t5 +-- !query analysis +Project [left(utf8_binary#x, cast(collate(2, utf8_lcase) as int)) AS left(utf8_binary, collate(2, utf8_lcase))#x, left(utf8_lcase#x, 2) AS left(utf8_lcase, 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select left(utf8_binary collate utf8_lcase, 2), left(utf8_lcase collate utf8_binary, 2) from t5 +-- !query analysis +Project [left(collate(utf8_binary#x, utf8_lcase), 2) AS left(collate(utf8_binary, utf8_lcase), 2)#x, left(collate(utf8_lcase#x, utf8_binary), 2) AS left(collate(utf8_lcase, utf8_binary), 2)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select rpad(utf8_binary, 8, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select rpad(s, 8, utf8_binary) from t5 +-- !query analysis +Project [rpad(s#x, 8, utf8_binary#x) AS rpad(s, 8, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select rpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select rpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [rpad(utf8_binary#x, 8, collate(utf8_lcase#x, utf8_binary)) AS rpad(utf8_binary, 8, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select rpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [rpad(collate(utf8_binary#x, utf8_lcase), 8, collate(utf8_lcase#x, utf8_lcase)) AS rpad(collate(utf8_binary, utf8_lcase), 8, collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select rpad(utf8_binary, 8, 'a'), rpad(utf8_lcase, 8, 'a') from t5 +-- !query analysis +Project [rpad(utf8_binary#x, 8, a) AS rpad(utf8_binary, 8, a)#x, rpad(utf8_lcase#x, 8, cast(a as string collate UTF8_LCASE)) AS rpad(utf8_lcase, 8, a)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select rpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), rpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5 +-- !query analysis +Project [rpad(cast(utf8_binary#x as string collate UTF8_LCASE), 8, collate(AaAA, utf8_lcase)) AS rpad(utf8_binary, 8, collate(AaAA, utf8_lcase))#x, rpad(cast(utf8_lcase#x as string), 8, collate(AAa, utf8_binary)) AS rpad(utf8_lcase, 8, collate(AAa, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select lpad(utf8_binary, 8, utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -468,7 +1940,15 @@ org.apache.spark.sql.AnalysisException -- !query -select split_part(str collate utf8_binary, delimiter collate utf8_lcase, partNum) from t5 +select lpad(s, 8, utf8_binary) from t5 +-- !query analysis +Project [lpad(s#x, 8, utf8_binary#x) AS lpad(s, 8, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select lpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -481,36 +1961,39 @@ org.apache.spark.sql.AnalysisException -- !query -select split_part(str collate utf8_binary, delimiter collate utf8_binary, partNum) from t5 +select lpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5 -- !query analysis -Project [split_part(collate(str#x, utf8_binary), collate(delimiter#x, utf8_binary), partNum#x) AS split_part(collate(str, utf8_binary), collate(delimiter, utf8_binary), partNum)#x] +Project [lpad(utf8_binary#x, 8, collate(utf8_lcase#x, utf8_binary)) AS lpad(utf8_binary, 8, collate(utf8_lcase, utf8_binary))#x] +- SubqueryAlias spark_catalog.default.t5 - +- Relation spark_catalog.default.t5[str#x,delimiter#x,partNum#x] parquet + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet -- !query -drop table t5 +select lpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5 -- !query analysis -DropTable false, false -+- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t5 +Project [lpad(collate(utf8_binary#x, utf8_lcase), 8, collate(utf8_lcase#x, utf8_lcase)) AS lpad(collate(utf8_binary, utf8_lcase), 8, collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet -- !query -create table t6 (utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase, threshold int) using parquet +select lpad(utf8_binary, 8, 'a'), lpad(utf8_lcase, 8, 'a') from t5 -- !query analysis -CreateDataSourceTableCommand `spark_catalog`.`default`.`t6`, false +Project [lpad(utf8_binary#x, 8, a) AS lpad(utf8_binary, 8, a)#x, lpad(utf8_lcase#x, 8, cast(a as string collate UTF8_LCASE)) AS lpad(utf8_lcase, 8, a)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet -- !query -insert into t6 values('kitten', 'sitting', 2) +select lpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), lpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5 -- !query analysis -InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t6, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/t6], Append, `spark_catalog`.`default`.`t6`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t6), [utf8_binary, utf8_lcase, threshold] -+- Project [cast(col1#x as string) AS utf8_binary#x, cast(col2#x as string collate UTF8_LCASE) AS utf8_lcase#x, cast(col3#x as int) AS threshold#x] - +- LocalRelation [col1#x, col2#x, col3#x] +Project [lpad(cast(utf8_binary#x as string collate UTF8_LCASE), 8, collate(AaAA, utf8_lcase)) AS lpad(utf8_binary, 8, collate(AaAA, utf8_lcase))#x, lpad(cast(utf8_lcase#x as string), 8, collate(AAa, utf8_binary)) AS lpad(utf8_lcase, 8, collate(AAa, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet -- !query -select levenshtein(utf8_binary, utf8_lcase) from t6 +select locate(utf8_binary, utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -520,7 +2003,15 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t6 +select locate(s, utf8_binary) from t5 +-- !query analysis +Project [locate(s#x, utf8_binary#x, 1) AS locate(s, utf8_binary, 1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select locate(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -533,15 +2024,102 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t6 +select locate(utf8_binary, utf8_lcase collate utf8_binary) from t5 -- !query analysis -Project [levenshtein(collate(utf8_binary#x, utf8_binary), collate(utf8_lcase#x, utf8_binary), None) AS levenshtein(collate(utf8_binary, utf8_binary), collate(utf8_lcase, utf8_binary))#x] -+- SubqueryAlias spark_catalog.default.t6 - +- Relation spark_catalog.default.t6[utf8_binary#x,utf8_lcase#x,threshold#x] parquet +Project [locate(utf8_binary#x, collate(utf8_lcase#x, utf8_binary), 1) AS locate(utf8_binary, collate(utf8_lcase, utf8_binary), 1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select locate(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 3) from t5 +-- !query analysis +Project [locate(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase), 3) AS locate(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase), 3)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5 +-- !query analysis +Project [locate(utf8_binary#x, a, 1) AS locate(utf8_binary, a, 1)#x, locate(utf8_lcase#x, cast(a as string collate UTF8_LCASE), 1) AS locate(utf8_lcase, a, 1)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select locate(utf8_binary, 'AaAA' collate utf8_lcase, 4), locate(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5 +-- !query analysis +Project [locate(cast(utf8_binary#x as string collate UTF8_LCASE), collate(AaAA, utf8_lcase), 4) AS locate(utf8_binary, collate(AaAA, utf8_lcase), 4)#x, locate(cast(utf8_lcase#x as string), collate(AAa, utf8_binary), 4) AS locate(utf8_lcase, collate(AAa, utf8_binary), 4)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select TRIM(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select TRIM(s, utf8_binary) from t5 +-- !query analysis +Project [trim(utf8_binary#x, Some(s#x)) AS TRIM(BOTH s FROM utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select TRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" + } +} + + +-- !query +select TRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [trim(collate(utf8_lcase#x, utf8_binary), Some(utf8_binary#x)) AS TRIM(BOTH utf8_binary FROM collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select TRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [trim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, utf8_lcase))) AS TRIM(BOTH collate(utf8_binary, utf8_lcase) FROM collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5 +-- !query analysis +Project [trim(utf8_binary#x, Some(ABc)) AS TRIM(BOTH ABc FROM utf8_binary)#x, trim(utf8_lcase#x, Some(cast(ABc as string collate UTF8_LCASE))) AS TRIM(BOTH ABc FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet -- !query -select levenshtein(utf8_binary, utf8_lcase, threshold) from t6 +select TRIM('ABc' collate utf8_lcase, utf8_binary), TRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query analysis +Project [trim(cast(utf8_binary#x as string collate UTF8_LCASE), Some(collate(ABc, utf8_lcase))) AS TRIM(BOTH collate(ABc, utf8_lcase) FROM utf8_binary)#x, trim(cast(utf8_lcase#x as string), Some(collate(AAa, utf8_binary))) AS TRIM(BOTH collate(AAa, utf8_binary) FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select BTRIM(utf8_binary, utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -551,7 +2129,15 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase, threshold) from t6 +select BTRIM(s, utf8_binary) from t5 +-- !query analysis +Project [btrim(s#x, utf8_binary#x) AS btrim(s, utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select BTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query analysis org.apache.spark.sql.AnalysisException { @@ -564,11 +2150,168 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary, threshold) from t6 +select BTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 -- !query analysis -Project [levenshtein(collate(utf8_binary#x, utf8_binary), collate(utf8_lcase#x, utf8_binary), Some(threshold#x)) AS levenshtein(collate(utf8_binary, utf8_binary), collate(utf8_lcase, utf8_binary), threshold)#x] -+- SubqueryAlias spark_catalog.default.t6 - +- Relation spark_catalog.default.t6[utf8_binary#x,utf8_lcase#x,threshold#x] parquet +Project [btrim(utf8_binary#x, collate(utf8_lcase#x, utf8_binary)) AS btrim(utf8_binary, collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select BTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [btrim(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lcase)) AS btrim(collate(utf8_binary, utf8_lcase), collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5 +-- !query analysis +Project [btrim(ABc, utf8_binary#x) AS btrim(ABc, utf8_binary)#x, btrim(ABc, utf8_lcase#x) AS btrim(ABc, utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select BTRIM('ABc' collate utf8_lcase, utf8_binary), BTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query analysis +Project [btrim(collate(ABc, utf8_lcase), utf8_binary#x) AS btrim(collate(ABc, utf8_lcase), utf8_binary)#x, btrim(collate(AAa, utf8_binary), utf8_lcase#x) AS btrim(collate(AAa, utf8_binary), utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select LTRIM(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select LTRIM(s, utf8_binary) from t5 +-- !query analysis +Project [ltrim(utf8_binary#x, Some(s#x)) AS TRIM(LEADING s FROM utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select LTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" + } +} + + +-- !query +select LTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [ltrim(collate(utf8_lcase#x, utf8_binary), Some(utf8_binary#x)) AS TRIM(LEADING utf8_binary FROM collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select LTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [ltrim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, utf8_lcase))) AS TRIM(LEADING collate(utf8_binary, utf8_lcase) FROM collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5 +-- !query analysis +Project [ltrim(utf8_binary#x, Some(ABc)) AS TRIM(LEADING ABc FROM utf8_binary)#x, ltrim(utf8_lcase#x, Some(cast(ABc as string collate UTF8_LCASE))) AS TRIM(LEADING ABc FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select LTRIM('ABc' collate utf8_lcase, utf8_binary), LTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query analysis +Project [ltrim(cast(utf8_binary#x as string collate UTF8_LCASE), Some(collate(ABc, utf8_lcase))) AS TRIM(LEADING collate(ABc, utf8_lcase) FROM utf8_binary)#x, ltrim(cast(utf8_lcase#x as string), Some(collate(AAa, utf8_binary))) AS TRIM(LEADING collate(AAa, utf8_binary) FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select RTRIM(utf8_binary, utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select RTRIM(s, utf8_binary) from t5 +-- !query analysis +Project [rtrim(utf8_binary#x, Some(s#x)) AS TRIM(TRAILING s FROM utf8_binary)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select RTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" + } +} + + +-- !query +select RTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query analysis +Project [rtrim(collate(utf8_lcase#x, utf8_binary), Some(utf8_binary#x)) AS TRIM(TRAILING utf8_binary FROM collate(utf8_lcase, utf8_binary))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select RTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query analysis +Project [rtrim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, utf8_lcase))) AS TRIM(TRAILING collate(utf8_binary, utf8_lcase) FROM collate(utf8_lcase, utf8_lcase))#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5 +-- !query analysis +Project [rtrim(utf8_binary#x, Some(ABc)) AS TRIM(TRAILING ABc FROM utf8_binary)#x, rtrim(utf8_lcase#x, Some(cast(ABc as string collate UTF8_LCASE))) AS TRIM(TRAILING ABc FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +select RTRIM('ABc' collate utf8_lcase, utf8_binary), RTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query analysis +Project [rtrim(cast(utf8_binary#x as string collate UTF8_LCASE), Some(collate(ABc, utf8_lcase))) AS TRIM(TRAILING collate(ABc, utf8_lcase) FROM utf8_binary)#x, rtrim(cast(utf8_lcase#x as string), Some(collate(AAa, utf8_binary))) AS TRIM(TRAILING collate(AAa, utf8_binary) FROM utf8_lcase)#x] ++- SubqueryAlias spark_catalog.default.t5 + +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet + + +-- !query +drop table t5 +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t5 -- !query @@ -576,3 +2319,31 @@ drop table t6 -- !query analysis DropTable false, false +- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t6 + + +-- !query +drop table t7 +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t7 + + +-- !query +drop table t8 +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t8 + + +-- !query +drop table t9 +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t9 + + +-- !query +drop table t10 +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t10 diff --git a/sql/core/src/test/resources/sql-tests/inputs/collations.sql b/sql/core/src/test/resources/sql-tests/inputs/collations.sql index 51d8d1be4154c..183577b83971b 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/collations.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/collations.sql @@ -102,27 +102,307 @@ select str_to_map(text collate utf8_binary, pairDelim collate utf8_binary, keyVa drop table t4; --- create table for split_part -create table t5(str string collate utf8_binary, delimiter string collate utf8_lcase, partNum int) using parquet; - -insert into t5 values('11AB12AB13', 'AB', 2); - -select split_part(str, delimiter, partNum) from t5; -select split_part(str collate utf8_binary, delimiter collate utf8_lcase, partNum) from t5; -select split_part(str collate utf8_binary, delimiter collate utf8_binary, partNum) from t5; +create table t5(s string, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet; +insert into t5 values ('Spark', 'Spark', 'SQL'); +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaAAaA'); +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaA'); +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaAaaAaaAaAaaAaaAaA'); +insert into t5 values ('bbAbaAbA', 'bbAbAAbA', 'a'); +insert into t5 values ('İo', 'İo', 'İo'); +insert into t5 values ('İo', 'İo', 'i̇o'); +insert into t5 values ('efd2', 'efd2', 'efd2'); +insert into t5 values ('Hello, world! Nice day.', 'Hello, world! Nice day.', 'Hello, world! Nice day.'); +insert into t5 values ('Something else. Nothing here.', 'Something else. Nothing here.', 'Something else. Nothing here.'); +insert into t5 values ('kitten', 'kitten', 'sitTing'); +insert into t5 values ('abc', 'abc', 'abc'); +insert into t5 values ('abcdcba', 'abcdcba', 'aBcDCbA'); + +create table t6(ascii long) using parquet; +insert into t6 values (97); +insert into t6 values (66); + +create table t7(ascii double) using parquet; +insert into t7 values (97.52143); +insert into t7 values (66.421); + +create table t8(format string collate utf8_binary, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet; +insert into t8 values ('%s%s', 'abCdE', 'abCdE'); + +create table t9(num long) using parquet; +insert into t9 values (97); +insert into t9 values (66); + +create table t10(utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet; +insert into t10 values ('aaAaAAaA', 'aaAaaAaA'); +insert into t10 values ('efd2', 'efd2'); + +-- ConcatWs +select concat_ws(' ', utf8_lcase, utf8_lcase) from t5; +select concat_ws(' ', utf8_binary, utf8_lcase) from t5; +select concat_ws(' ' collate utf8_binary, utf8_binary, 'SQL' collate utf8_lcase) from t5; +select concat_ws(' ' collate utf8_lcase, utf8_binary, 'SQL' collate utf8_lcase) from t5; +select concat_ws(',', utf8_lcase, 'word'), concat_ws(',', utf8_binary, 'word') from t5; +select concat_ws(',', utf8_lcase, 'word' collate utf8_binary), concat_ws(',', utf8_binary, 'word' collate utf8_lcase) from t5; + +-- Elt +select elt(2, s, utf8_binary) from t5; +select elt(2, utf8_binary, utf8_lcase, s) from t5; +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t5; +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t5; +select elt(1, utf8_binary collate utf8_binary, utf8_lcase) from t5; +select elt(1, utf8_binary, 'word'), elt(1, utf8_lcase, 'word') from t5; +select elt(1, utf8_binary, 'word' collate utf8_lcase), elt(1, utf8_lcase, 'word' collate utf8_binary) from t5; + +-- SplitPart +select split_part(utf8_binary, utf8_lcase, 3) from t5; +select split_part(s, utf8_binary, 1) from t5; +select split_part(utf8_binary collate utf8_binary, s collate utf8_lcase, 1) from t5; +select split_part(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5; +select split_part(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5; +select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5; +select split_part(utf8_binary, 'a' collate utf8_lcase, 3), split_part(utf8_lcase, 'a' collate utf8_binary, 3) from t5; + +-- Contains +select contains(utf8_binary, utf8_lcase) from t5; +select contains(s, utf8_binary) from t5; +select contains(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select contains(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select contains(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5; +select contains(utf8_binary, 'AaAA' collate utf8_lcase), contains(utf8_lcase, 'AAa' collate utf8_binary) from t5; + +-- SubstringIndex +select substring_index(utf8_binary, utf8_lcase, 2) from t5; +select substring_index(s, utf8_binary,1) from t5; +select substring_index(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5; +select substring_index(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5; +select substring_index(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5; +select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5; +select substring_index(utf8_binary, 'AaAA' collate utf8_lcase, 2), substring_index(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5; + +-- StringInStr +select instr(utf8_binary, utf8_lcase) from t5; +select instr(s, utf8_binary) from t5; +select instr(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select instr(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select instr(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5; +select instr(utf8_binary, 'AaAA' collate utf8_lcase), instr(utf8_lcase, 'AAa' collate utf8_binary) from t5; + +-- FindInSet +select find_in_set(utf8_binary, utf8_lcase) from t5; +select find_in_set(s, utf8_binary) from t5; +select find_in_set(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select find_in_set(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select find_in_set(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o'), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o') from t5; +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o' collate utf8_lcase), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o' collate utf8_binary) from t5; + +-- StartsWith +select startswith(utf8_binary, utf8_lcase) from t5; +select startswith(s, utf8_binary) from t5; +select startswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select startswith(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select startswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5; +select startswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), startswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5; + +-- StringTranslate +select translate(utf8_lcase, utf8_lcase, '12345') from t5; +select translate(utf8_binary, utf8_lcase, '12345') from t5; +select translate(utf8_binary, 'aBc' collate utf8_lcase, '12345' collate utf8_binary) from t5; +select translate(utf8_binary, 'SQL' collate utf8_lcase, '12345' collate utf8_lcase) from t5; +select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5; +select translate(utf8_lcase, 'aBc' collate utf8_binary, '12345'), translate(utf8_binary, 'aBc' collate utf8_lcase, '12345') from t5; + +-- Replace +select replace(utf8_binary, utf8_lcase, 'abc') from t5; +select replace(s, utf8_binary, 'abc') from t5; +select replace(utf8_binary collate utf8_binary, s collate utf8_lcase, 'abc') from t5; +select replace(utf8_binary, utf8_lcase collate utf8_binary, 'abc') from t5; +select replace(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 'abc') from t5; +select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5; +select replace(utf8_binary, 'aaAaaAaA' collate utf8_lcase, 'abc'), replace(utf8_lcase, 'aaAaaAaA' collate utf8_binary, 'abc') from t5; + +-- EndsWith +select endswith(utf8_binary, utf8_lcase) from t5; +select endswith(s, utf8_binary) from t5; +select endswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select endswith(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select endswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5; +select endswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), endswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5; + +-- StringRepeat +select repeat(utf8_binary, 3), repeat(utf8_lcase, 2) from t5; +select repeat(utf8_binary collate utf8_lcase, 3), repeat(utf8_lcase collate utf8_binary, 2) from t5; + +-- Ascii & UnBase64 string expressions +select ascii(utf8_binary), ascii(utf8_lcase) from t5; +select ascii(utf8_binary collate utf8_lcase), ascii(utf8_lcase collate utf8_binary) from t5; +select unbase64(utf8_binary), unbase64(utf8_lcase) from t10; +select unbase64(utf8_binary collate utf8_lcase), unbase64(utf8_lcase collate utf8_binary) from t10; + +-- Chr +select chr(ascii) from t6; + +-- Base64, Decode +select base64(utf8_binary), base64(utf8_lcase) from t5; +select base64(utf8_binary collate utf8_lcase), base64(utf8_lcase collate utf8_binary) from t5; +select decode(encode(utf8_binary, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase, 'utf-8'), 'utf-8') from t5; +select decode(encode(utf8_binary collate utf8_lcase, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase collate utf8_binary, 'utf-8'), 'utf-8') from t5; + +-- FormatNumber +select format_number(ascii, '###.###') from t7; +select format_number(ascii, '###.###' collate utf8_lcase) from t7; + +-- Encode, ToBinary +select encode(utf8_binary, 'utf-8'), encode(utf8_lcase, 'utf-8') from t5; +select encode(utf8_binary collate utf8_lcase, 'utf-8'), encode(utf8_lcase collate utf8_binary, 'utf-8') from t5; +select to_binary(utf8_binary, 'utf-8'), to_binary(utf8_lcase, 'utf-8') from t5; +select to_binary(utf8_binary collate utf8_lcase, 'utf-8'), to_binary(utf8_lcase collate utf8_binary, 'utf-8') from t5; + +-- Sentences +select sentences(utf8_binary), sentences(utf8_lcase) from t5; +select sentences(utf8_binary collate utf8_lcase), sentences(utf8_lcase collate utf8_binary) from t5; + +-- Upper +select upper(utf8_binary), upper(utf8_lcase) from t5; +select upper(utf8_binary collate utf8_lcase), upper(utf8_lcase collate utf8_binary) from t5; + +-- Lower +select lower(utf8_binary), lower(utf8_lcase) from t5; +select lower(utf8_binary collate utf8_lcase), lower(utf8_lcase collate utf8_binary) from t5; + +-- InitCap +select initcap(utf8_binary), initcap(utf8_lcase) from t5; +select initcap(utf8_binary collate utf8_lcase), initcap(utf8_lcase collate utf8_binary) from t5; + +-- Overlay +select overlay(utf8_binary, utf8_lcase, 2) from t5; +select overlay(s, utf8_binary,1) from t5; +select overlay(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5; +select overlay(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5; +select overlay(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5; +select overlay(utf8_binary, 'a', 2), overlay(utf8_lcase, 'a', 2) from t5; +select overlay(utf8_binary, 'AaAA' collate utf8_lcase, 2), overlay(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5; + +-- FormatString +select format_string(format, utf8_binary, utf8_lcase) from t8; +select format_string(format collate utf8_lcase, utf8_lcase, utf8_binary collate utf8_lcase, 3), format_string(format, utf8_lcase collate utf8_binary, utf8_binary) from t8; +select format_string(format, utf8_binary, utf8_lcase) from t8; + +-- SoundEx +select soundex(utf8_binary), soundex(utf8_lcase) from t5; +select soundex(utf8_binary collate utf8_lcase), soundex(utf8_lcase collate utf8_binary) from t5; + +-- Length, BitLength & OctetLength +select length(utf8_binary), length(utf8_lcase) from t5; +select length(utf8_binary collate utf8_lcase), length(utf8_lcase collate utf8_binary) from t5; +select bit_length(utf8_binary), bit_length(utf8_lcase) from t5; +select bit_length(utf8_binary collate utf8_lcase), bit_length(utf8_lcase collate utf8_binary) from t5; +select octet_length(utf8_binary), octet_length(utf8_lcase) from t5; +select octet_length(utf8_binary collate utf8_lcase), octet_length(utf8_lcase collate utf8_binary) from t5; + +-- Luhncheck +select luhn_check(num) from t9; + +-- Levenshtein +select levenshtein(utf8_binary, utf8_lcase) from t5; +select levenshtein(s, utf8_binary) from t5; +select levenshtein(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select levenshtein(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select levenshtein(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select levenshtein(utf8_binary, 'a'), levenshtein(utf8_lcase, 'a') from t5; +select levenshtein(utf8_binary, 'AaAA' collate utf8_lcase, 3), levenshtein(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5; + +-- IsValidUTF8 +select is_valid_utf8(utf8_binary), is_valid_utf8(utf8_lcase) from t5; +select is_valid_utf8(utf8_binary collate utf8_lcase), is_valid_utf8(utf8_lcase collate utf8_binary) from t5; + +-- MakeValidUTF8 +select make_valid_utf8(utf8_binary), make_valid_utf8(utf8_lcase) from t5; +select make_valid_utf8(utf8_binary collate utf8_lcase), make_valid_utf8(utf8_lcase collate utf8_binary) from t5; + +-- ValidateUTF8 +select validate_utf8(utf8_binary), validate_utf8(utf8_lcase) from t5; +select validate_utf8(utf8_binary collate utf8_lcase), validate_utf8(utf8_lcase collate utf8_binary) from t5; + +-- TryValidateUTF8 +select try_validate_utf8(utf8_binary), try_validate_utf8(utf8_lcase) from t5; +select try_validate_utf8(utf8_binary collate utf8_lcase), try_validate_utf8(utf8_lcase collate utf8_binary) from t5; + +-- Left/Right/Substr +select substr(utf8_binary, 2, 2), substr(utf8_lcase, 2, 2) from t5; +select substr(utf8_binary collate utf8_lcase, 2, 2), substr(utf8_lcase collate utf8_binary, 2, 2) from t5; +select right(utf8_binary, 2), right(utf8_lcase, 2) from t5; +select right(utf8_binary collate utf8_lcase, 2), right(utf8_lcase collate utf8_binary, 2) from t5; +select left(utf8_binary, '2' collate utf8_lcase), left(utf8_lcase, 2) from t5; +select left(utf8_binary collate utf8_lcase, 2), left(utf8_lcase collate utf8_binary, 2) from t5; + +-- StringRPad +select rpad(utf8_binary, 8, utf8_lcase) from t5; +select rpad(s, 8, utf8_binary) from t5; +select rpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5; +select rpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5; +select rpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5; +select rpad(utf8_binary, 8, 'a'), rpad(utf8_lcase, 8, 'a') from t5; +select rpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), rpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5; + +-- StringLPad +select lpad(utf8_binary, 8, utf8_lcase) from t5; +select lpad(s, 8, utf8_binary) from t5; +select lpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5; +select lpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5; +select lpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5; +select lpad(utf8_binary, 8, 'a'), lpad(utf8_lcase, 8, 'a') from t5; +select lpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), lpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5; + +-- Locate +select locate(utf8_binary, utf8_lcase) from t5; +select locate(s, utf8_binary) from t5; +select locate(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select locate(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select locate(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 3) from t5; +select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5; +select locate(utf8_binary, 'AaAA' collate utf8_lcase, 4), locate(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5; + +-- StringTrim +select TRIM(utf8_binary, utf8_lcase) from t5; +select TRIM(s, utf8_binary) from t5; +select TRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select TRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select TRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5; +select TRIM('ABc' collate utf8_lcase, utf8_binary), TRIM('AAa' collate utf8_binary, utf8_lcase) from t5; +-- StringTrimBoth +select BTRIM(utf8_binary, utf8_lcase) from t5; +select BTRIM(s, utf8_binary) from t5; +select BTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select BTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select BTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5; +select BTRIM('ABc' collate utf8_lcase, utf8_binary), BTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; +-- StringTrimLeft +select LTRIM(utf8_binary, utf8_lcase) from t5; +select LTRIM(s, utf8_binary) from t5; +select LTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select LTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select LTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5; +select LTRIM('ABc' collate utf8_lcase, utf8_binary), LTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; +-- StringTrimRight +select RTRIM(utf8_binary, utf8_lcase) from t5; +select RTRIM(s, utf8_binary) from t5; +select RTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; +select RTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; +select RTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5; +select RTRIM('ABc' collate utf8_lcase, utf8_binary), RTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; drop table t5; - --- create table for levenshtein -create table t6 (utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase, threshold int) using parquet; - -insert into t6 values('kitten', 'sitting', 2); - -select levenshtein(utf8_binary, utf8_lcase) from t6; -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t6; -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t6; -select levenshtein(utf8_binary, utf8_lcase, threshold) from t6; -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase, threshold) from t6; -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary, threshold) from t6; - drop table t6; +drop table t7; +drop table t8; +drop table t9; +drop table t10; diff --git a/sql/core/src/test/resources/sql-tests/results/collations.sql.out b/sql/core/src/test/resources/sql-tests/results/collations.sql.out index d8f8d0676baed..ea5564aafe96f 100644 --- a/sql/core/src/test/resources/sql-tests/results/collations.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/collations.sql.out @@ -489,7 +489,7 @@ struct<> -- !query -create table t5(str string collate utf8_binary, delimiter string collate utf8_lcase, partNum int) using parquet +create table t5(s string, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet -- !query schema struct<> -- !query output @@ -497,7 +497,7 @@ struct<> -- !query -insert into t5 values('11AB12AB13', 'AB', 2) +insert into t5 values ('Spark', 'Spark', 'SQL') -- !query schema struct<> -- !query output @@ -505,7 +505,3113 @@ struct<> -- !query -select split_part(str, delimiter, partNum) from t5 +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaAAaA') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaA') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('aaAaAAaA', 'aaAaAAaA', 'aaAaaAaAaaAaaAaAaaAaaAaA') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('bbAbaAbA', 'bbAbAAbA', 'a') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('İo', 'İo', 'İo') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('İo', 'İo', 'i̇o') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('efd2', 'efd2', 'efd2') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('Hello, world! Nice day.', 'Hello, world! Nice day.', 'Hello, world! Nice day.') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('Something else. Nothing here.', 'Something else. Nothing here.', 'Something else. Nothing here.') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('kitten', 'kitten', 'sitTing') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('abc', 'abc', 'abc') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t5 values ('abcdcba', 'abcdcba', 'aBcDCbA') +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t6(ascii long) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t6 values (97) +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t6 values (66) +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t7(ascii double) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t7 values (97.52143) +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t7 values (66.421) +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t8(format string collate utf8_binary, utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t8 values ('%s%s', 'abCdE', 'abCdE') +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t9(num long) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t9 values (97) +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t9 values (66) +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t10(utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t10 values ('aaAaAAaA', 'aaAaaAaA') +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t10 values ('efd2', 'efd2') +-- !query schema +struct<> +-- !query output + + + +-- !query +select concat_ws(' ', utf8_lcase, utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +SQL SQL +Something else. Nothing here. Something else. Nothing here. +a a +aBcDCbA aBcDCbA +aaAaAAaA aaAaAAaA +aaAaaAaA aaAaaAaA +aaAaaAaAaaAaaAaAaaAaaAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +efd2 efd2 +i̇o i̇o +sitTing sitTing +İo İo + + +-- !query +select concat_ws(' ', utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select concat_ws(' ' collate utf8_binary, utf8_binary, 'SQL' collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select concat_ws(' ' collate utf8_lcase, utf8_binary, 'SQL' collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. SQL +Something else. Nothing here. SQL +Spark SQL +aaAaAAaA SQL +aaAaAAaA SQL +aaAaAAaA SQL +abc SQL +abcdcba SQL +bbAbAAbA SQL +efd2 SQL +kitten SQL +İo SQL +İo SQL + + +-- !query +select concat_ws(',', utf8_lcase, 'word'), concat_ws(',', utf8_binary, 'word') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day.,word Hello, world! Nice day.,word +SQL,word Spark,word +Something else. Nothing here.,word Something else. Nothing here.,word +a,word bbAbAAbA,word +aBcDCbA,word abcdcba,word +aaAaAAaA,word aaAaAAaA,word +aaAaaAaA,word aaAaAAaA,word +aaAaaAaAaaAaaAaAaaAaaAaA,word aaAaAAaA,word +abc,word abc,word +efd2,word efd2,word +i̇o,word İo,word +sitTing,word kitten,word +İo,word İo,word + + +-- !query +select concat_ws(',', utf8_lcase, 'word' collate utf8_binary), concat_ws(',', utf8_binary, 'word' collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day.,word Hello, world! Nice day.,word +SQL,word Spark,word +Something else. Nothing here.,word Something else. Nothing here.,word +a,word bbAbAAbA,word +aBcDCbA,word abcdcba,word +aaAaAAaA,word aaAaAAaA,word +aaAaaAaA,word aaAaAAaA,word +aaAaaAaAaaAaaAaAaaAaaAaA,word aaAaAAaA,word +abc,word abc,word +efd2,word efd2,word +i̇o,word İo,word +sitTing,word kitten,word +İo,word İo,word + + +-- !query +select elt(2, s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select elt(2, utf8_binary, utf8_lcase, s) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select elt(1, utf8_binary collate utf8_binary, utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select elt(1, utf8_binary, 'word'), elt(1, utf8_lcase, 'word') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select elt(1, utf8_binary, 'word' collate utf8_lcase), elt(1, utf8_lcase, 'word' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select split_part(utf8_binary, utf8_lcase, 3) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select split_part(s, utf8_binary, 1) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + + +bbAbaAbA + + +-- !query +select split_part(utf8_binary collate utf8_binary, s collate utf8_lcase, 1) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select split_part(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output + + + +-- !query +select split_part(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + + +b + + +-- !query +select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + +A +A +A + + +-- !query +select split_part(utf8_binary, 'a' collate utf8_lcase, 3), split_part(utf8_lcase, 'a' collate utf8_binary, 3) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + A + A + A + + +-- !query +select contains(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select contains(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +true +true +true +true +true +true +true +true +true +true +true +true + + +-- !query +select contains(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select contains(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +false +false +false +false +false +false +true +true +true +true +true +true + + +-- !query +select contains(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +false +false +false +true +true +true +true +true +true +true +true +true +true + + +-- !query +select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false true +true false +true true +true true +true true +true true +true true +true true + + +-- !query +select contains(utf8_binary, 'AaAA' collate utf8_lcase), contains(utf8_lcase, 'AAa' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false false +false false +false false +false false +false false +true false +true false +true true + + +-- !query +select substring_index(utf8_binary, utf8_lcase, 2) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select substring_index(s, utf8_binary,1) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + + +bbAbaAbA + + +-- !query +select substring_index(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select substring_index(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select substring_index(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAb +efd2 +kitten +İo +İo + + +-- !query +select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +a a +a a +a a +abc abc +abcdcb aBcDCb +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select substring_index(utf8_binary, 'AaAA' collate utf8_lcase, 2), substring_index(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +a aaAaAAaA +a aaAaaAaA +a aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select instr(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select instr(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 + + +-- !query +select instr(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select instr(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 +1 + + +-- !query +select instr(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +3 + + +-- !query +select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 1 +1 1 +1 1 +1 1 +1 1 +1 1 +21 21 +3 0 + + +-- !query +select instr(utf8_binary, 'AaAA' collate utf8_lcase), instr(utf8_lcase, 'AAa' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +1 0 +1 0 +1 5 + + +-- !query +select find_in_set(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select find_in_set(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 + + +-- !query +select find_in_set(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select find_in_set(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 + + +-- !query +select find_in_set(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 +1 +1 +1 + + +-- !query +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o'), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o') from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 1 +0 1 +0 2 +0 2 + + +-- !query +select find_in_set(utf8_binary, 'aaAaaAaA,i̇o' collate utf8_lcase), find_in_set(utf8_lcase, 'aaAaaAaA,i̇o' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +1 0 +1 0 +1 1 +2 0 +2 2 + + +-- !query +select startswith(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select startswith(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +true +true +true +true +true +true +true +true +true +true +true +true + + +-- !query +select startswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select startswith(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +false +false +false +false +false +false +true +true +true +true +true +true + + +-- !query +select startswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +false +false +false +false +true +true +true +true +true +true +true +true +true + + +-- !query +select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false false +false false +false false +false false +false false +false true +false true +false true + + +-- !query +select startswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), startswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false false +false false +false false +false false +false false +true false +true true +true true + + +-- !query +select translate(utf8_lcase, utf8_lcase, '12345') from t5 +-- !query schema +struct +-- !query output +1 +11111111 +11111111 +111111111111111111111111 +12 +123 +123 +123 +12332 +12335532 +1234 +1234321 +123454142544 + + +-- !query +select translate(utf8_binary, utf8_lcase, '12345') from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select translate(utf8_binary, 'aBc' collate utf8_lcase, '12345' collate utf8_binary) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" + } +} + + +-- !query +select translate(utf8_binary, 'SQL' collate utf8_lcase, '12345' collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +1omething e31e. Nothing here. +1park +He33o, wor3d! Nice day. +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5 +-- !query schema +struct +-- !query output +1 bb3b33b3 +11111111 11313313 +11111111 11313313 +111111111111111111111111 11313313 +1BcDCb1 1bcdcb1 +1bc 1bc +Hello, world! Nice d1y. Hello, world! Nice d1y. +SQL Sp1rk +Something else. Nothing here. Something else. Nothing here. +efd2 efd2 +i̇o İo +sitTing kitten +İo İo + + +-- !query +select translate(utf8_lcase, 'aBc' collate utf8_binary, '12345'), translate(utf8_binary, 'aBc' collate utf8_lcase, '12345') from t5 +-- !query schema +struct +-- !query output +1 22121121 +11A11A1A 11111111 +11A11A1A11A11A1A11A11A1A 11111111 +11A1AA1A 11111111 +123DCbA 123d321 +1b3 123 +Hello, world! Ni3e d1y. Hello, world! Ni3e d1y. +SQL Sp1rk +Something else. Nothing here. Something else. Nothing here. +efd2 efd2 +i̇o İo +sitTing kitten +İo İo + + +-- !query +select replace(utf8_binary, utf8_lcase, 'abc') from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select replace(s, utf8_binary, 'abc') from t5 +-- !query schema +struct +-- !query output +abc +abc +abc +abc +abc +abc +abc +abc +abc +abc +abc +abc +bbAbaAbA + + +-- !query +select replace(utf8_binary collate utf8_binary, s collate utf8_lcase, 'abc') from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select replace(utf8_binary, utf8_lcase collate utf8_binary, 'abc') from t5 +-- !query schema +struct +-- !query output +Spark +aaAaAAaA +aaAaAAaA +abc +abc +abc +abc +abc +abc +abcdcba +bbAbAAbA +kitten +İo + + +-- !query +select replace(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 'abc') from t5 +-- !query schema +struct +-- !query output +Spark +aaAaAAaA +abc +abc +abc +abc +abc +abc +abc +abc +abc +bbabcbabcabcbabc +kitten + + +-- !query +select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA abc +aaAaAAaA abc +aaAaAAaA abcabcabc +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select replace(utf8_binary, 'aaAaaAaA' collate utf8_lcase, 'abc'), replace(utf8_lcase, 'aaAaaAaA' collate utf8_binary, 'abc') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +abc aaAaAAaA +abc abc +abc abc +abc abcabcabc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select endswith(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select endswith(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +true +true +true +true +true +true +true +true +true +true +true +true + + +-- !query +select endswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select endswith(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false +false +false +false +false +false +false +true +true +true +true +true +true + + +-- !query +select endswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +false +false +false +false +false +true +true +true +true +true +true +true +true + + +-- !query +select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false false +false false +false false +false false +false false +false true +false true +false true + + +-- !query +select endswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), endswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +false false +false false +false false +false false +false false +false false +false false +false false +false false +false false +true false +true true +true true + + +-- !query +select repeat(utf8_binary, 3), repeat(utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day.Hello, world! Nice day.Hello, world! Nice day. Hello, world! Nice day.Hello, world! Nice day. +Something else. Nothing here.Something else. Nothing here.Something else. Nothing here. Something else. Nothing here.Something else. Nothing here. +SparkSparkSpark SQLSQL +aaAaAAaAaaAaAAaAaaAaAAaA aaAaAAaAaaAaAAaA +aaAaAAaAaaAaAAaAaaAaAAaA aaAaaAaAaaAaaAaA +aaAaAAaAaaAaAAaAaaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaA +abcabcabc abcabc +abcdcbaabcdcbaabcdcba aBcDCbAaBcDCbA +bbAbAAbAbbAbAAbAbbAbAAbA aa +efd2efd2efd2 efd2efd2 +kittenkittenkitten sitTingsitTing +İoİoİo i̇oi̇o +İoİoİo İoİo + + +-- !query +select repeat(utf8_binary collate utf8_lcase, 3), repeat(utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day.Hello, world! Nice day.Hello, world! Nice day. Hello, world! Nice day.Hello, world! Nice day. +Something else. Nothing here.Something else. Nothing here.Something else. Nothing here. Something else. Nothing here.Something else. Nothing here. +SparkSparkSpark SQLSQL +aaAaAAaAaaAaAAaAaaAaAAaA aaAaAAaAaaAaAAaA +aaAaAAaAaaAaAAaAaaAaAAaA aaAaaAaAaaAaaAaA +aaAaAAaAaaAaAAaAaaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaAaaAaaAaA +abcabcabc abcabc +abcdcbaabcdcbaabcdcba aBcDCbAaBcDCbA +bbAbAAbAbbAbAAbAbbAbAAbA aa +efd2efd2efd2 efd2efd2 +kittenkittenkitten sitTingsitTing +İoİoİo i̇oi̇o +İoİoİo İoİo + + +-- !query +select ascii(utf8_binary), ascii(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +101 101 +107 115 +304 105 +304 304 +72 72 +83 83 +83 83 +97 97 +97 97 +97 97 +97 97 +97 97 +98 97 + + +-- !query +select ascii(utf8_binary collate utf8_lcase), ascii(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +101 101 +107 115 +304 105 +304 304 +72 72 +83 83 +83 83 +97 97 +97 97 +97 97 +97 97 +97 97 +98 97 + + +-- !query +select unbase64(utf8_binary), unbase64(utf8_lcase) from t10 +-- !query schema +struct +-- !query output +i�� i�h� +y�v y�v + + +-- !query +select unbase64(utf8_binary collate utf8_lcase), unbase64(utf8_lcase collate utf8_binary) from t10 +-- !query schema +struct +-- !query output +i�� i�h� +y�v y�v + + +-- !query +select chr(ascii) from t6 +-- !query schema +struct +-- !query output +B +a + + +-- !query +select base64(utf8_binary), base64(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +SGVsbG8sIHdvcmxkISBOaWNlIGRheS4= SGVsbG8sIHdvcmxkISBOaWNlIGRheS4= +U29tZXRoaW5nIGVsc2UuIE5vdGhpbmcgaGVyZS4= U29tZXRoaW5nIGVsc2UuIE5vdGhpbmcgaGVyZS4= +U3Bhcms= U1FM +YWFBYUFBYUE= YWFBYUFBYUE= +YWFBYUFBYUE= YWFBYWFBYUE= +YWFBYUFBYUE= YWFBYWFBYUFhYUFhYUFhQWFhQWFhQWFB +YWJj YWJj +YWJjZGNiYQ== YUJjRENiQQ== +YmJBYkFBYkE= YQ== +ZWZkMg== ZWZkMg== +a2l0dGVu c2l0VGluZw== +xLBv acyHbw== +xLBv xLBv + + +-- !query +select base64(utf8_binary collate utf8_lcase), base64(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +SGVsbG8sIHdvcmxkISBOaWNlIGRheS4= SGVsbG8sIHdvcmxkISBOaWNlIGRheS4= +U29tZXRoaW5nIGVsc2UuIE5vdGhpbmcgaGVyZS4= U29tZXRoaW5nIGVsc2UuIE5vdGhpbmcgaGVyZS4= +U3Bhcms= U1FM +YWFBYUFBYUE= YWFBYUFBYUE= +YWFBYUFBYUE= YWFBYWFBYUE= +YWFBYUFBYUE= YWFBYWFBYUFhYUFhYUFhQWFhQWFhQWFB +YWJj YWJj +YWJjZGNiYQ== YUJjRENiQQ== +YmJBYkFBYkE= YQ== +ZWZkMg== ZWZkMg== +a2l0dGVu c2l0VGluZw== +xLBv acyHbw== +xLBv xLBv + + +-- !query +select decode(encode(utf8_binary, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase, 'utf-8'), 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select decode(encode(utf8_binary collate utf8_lcase, 'utf-8'), 'utf-8'), decode(encode(utf8_lcase collate utf8_binary, 'utf-8'), 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select format_number(ascii, '###.###') from t7 +-- !query schema +struct +-- !query output +66.421 +97.521 + + +-- !query +select format_number(ascii, '###.###' collate utf8_lcase) from t7 +-- !query schema +struct +-- !query output +66.421 +97.521 + + +-- !query +select encode(utf8_binary, 'utf-8'), encode(utf8_lcase, 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select encode(utf8_binary collate utf8_lcase, 'utf-8'), encode(utf8_lcase collate utf8_binary, 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select to_binary(utf8_binary, 'utf-8'), to_binary(utf8_lcase, 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select to_binary(utf8_binary collate utf8_lcase, 'utf-8'), to_binary(utf8_lcase collate utf8_binary, 'utf-8') from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select sentences(utf8_binary), sentences(utf8_lcase) from t5 +-- !query schema +struct>,sentences(utf8_lcase, , ):array>> +-- !query output +[["Hello","world"],["Nice","day"]] [["Hello","world"],["Nice","day"]] +[["Something","else"],["Nothing","here"]] [["Something","else"],["Nothing","here"]] +[["Spark"]] [["SQL"]] +[["aaAaAAaA"]] [["aaAaAAaA"]] +[["aaAaAAaA"]] [["aaAaaAaA"]] +[["aaAaAAaA"]] [["aaAaaAaAaaAaaAaAaaAaaAaA"]] +[["abc"]] [["abc"]] +[["abcdcba"]] [["aBcDCbA"]] +[["bbAbAAbA"]] [["a"]] +[["efd2"]] [["efd2"]] +[["kitten"]] [["sitTing"]] +[["İo"]] [["i̇o"]] +[["İo"]] [["İo"]] + + +-- !query +select sentences(utf8_binary collate utf8_lcase), sentences(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct>,sentences(collate(utf8_lcase, utf8_binary), , ):array>> +-- !query output +[["Hello","world"],["Nice","day"]] [["Hello","world"],["Nice","day"]] +[["Something","else"],["Nothing","here"]] [["Something","else"],["Nothing","here"]] +[["Spark"]] [["SQL"]] +[["aaAaAAaA"]] [["aaAaAAaA"]] +[["aaAaAAaA"]] [["aaAaaAaA"]] +[["aaAaAAaA"]] [["aaAaaAaAaaAaaAaAaaAaaAaA"]] +[["abc"]] [["abc"]] +[["abcdcba"]] [["aBcDCbA"]] +[["bbAbAAbA"]] [["a"]] +[["efd2"]] [["efd2"]] +[["kitten"]] [["sitTing"]] +[["İo"]] [["i̇o"]] +[["İo"]] [["İo"]] + + +-- !query +select upper(utf8_binary), upper(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +AAAAAAAA AAAAAAAA +AAAAAAAA AAAAAAAA +AAAAAAAA AAAAAAAAAAAAAAAAAAAAAAAA +ABC ABC +ABCDCBA ABCDCBA +BBABAABA A +EFD2 EFD2 +HELLO, WORLD! NICE DAY. HELLO, WORLD! NICE DAY. +KITTEN SITTING +SOMETHING ELSE. NOTHING HERE. SOMETHING ELSE. NOTHING HERE. +SPARK SQL +İO İO +İO İO + + +-- !query +select upper(utf8_binary collate utf8_lcase), upper(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +AAAAAAAA AAAAAAAA +AAAAAAAA AAAAAAAA +AAAAAAAA AAAAAAAAAAAAAAAAAAAAAAAA +ABC ABC +ABCDCBA ABCDCBA +BBABAABA A +EFD2 EFD2 +HELLO, WORLD! NICE DAY. HELLO, WORLD! NICE DAY. +KITTEN SITTING +SOMETHING ELSE. NOTHING HERE. SOMETHING ELSE. NOTHING HERE. +SPARK SQL +İO İO +İO İO + + +-- !query +select lower(utf8_binary), lower(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +aaaaaaaa aaaaaaaa +aaaaaaaa aaaaaaaa +aaaaaaaa aaaaaaaaaaaaaaaaaaaaaaaa +abc abc +abcdcba abcdcba +bbabaaba a +efd2 efd2 +hello, world! nice day. hello, world! nice day. +i̇o i̇o +i̇o i̇o +kitten sitting +something else. nothing here. something else. nothing here. +spark sql + + +-- !query +select lower(utf8_binary collate utf8_lcase), lower(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +aaaaaaaa aaaaaaaa +aaaaaaaa aaaaaaaa +aaaaaaaa aaaaaaaaaaaaaaaaaaaaaaaa +abc abc +abcdcba abcdcba +bbabaaba a +efd2 efd2 +hello, world! nice day. hello, world! nice day. +i̇o i̇o +i̇o i̇o +kitten sitting +something else. nothing here. something else. nothing here. +spark sql + + +-- !query +select initcap(utf8_binary), initcap(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Aaaaaaaa Aaaaaaaa +Aaaaaaaa Aaaaaaaa +Aaaaaaaa Aaaaaaaaaaaaaaaaaaaaaaaa +Abc Abc +Abcdcba Abcdcba +Bbabaaba A +Efd2 Efd2 +Hello, World! Nice Day. Hello, World! Nice Day. +Kitten Sitting +Something Else. Nothing Here. Something Else. Nothing Here. +Spark Sql +İo İo +İo İo + + +-- !query +select initcap(utf8_binary collate utf8_lcase), initcap(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Aaaaaaaa Aaaaaaaa +Aaaaaaaa Aaaaaaaa +Aaaaaaaa Aaaaaaaaaaaaaaaaaaaaaaaa +Abc Abc +Abcdcba Abcdcba +Bbabaaba A +Efd2 Efd2 +Hello, World! Nice Day. Hello, World! Nice Day. +Kitten Sitting +Something Else. Nothing Here. Something Else. Nothing Here. +Spark Sql +İo İo +İo İo + + +-- !query +select overlay(utf8_binary, utf8_lcase, 2) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select overlay(s, utf8_binary,1) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. +Something else. Nothing here. +Spark +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba +bbAbAAbA +efd2 +kitten +İo +İo + + +-- !query +select overlay(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select overlay(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +HHello, world! Nice day. +SSQLk +SSomething else. Nothing here. +aaBcDCbA +aaaAaAAaA +aaaAaaAaA +aaaAaaAaAaaAaaAaAaaAaaAaA +aabc +baAbAAbA +eefd2 +ksitTing +İi̇o +İİo + + +-- !query +select overlay(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output +HHello, world! Nice day. +SSQLk +SSomething else. Nothing here. +aaBcDCbA +aaaAaAAaA +aaaAaaAaA +aaaAaaAaAaaAaaAaAaaAaaAaA +aabc +baAbAAbA +eefd2 +ksitTing +İi̇o +İİo + + +-- !query +select overlay(utf8_binary, 'a', 2), overlay(utf8_lcase, 'a', 2) from t5 +-- !query schema +struct +-- !query output +Hallo, world! Nice day. Hallo, world! Nice day. +Saark SaL +Samething else. Nothing here. Samething else. Nothing here. +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +aac aac +aacdcba aacDCbA +baAbAAbA aa +ead2 ead2 +katten satTing +İa iao +İa İa + + +-- !query +select overlay(utf8_binary, 'AaAA' collate utf8_lcase, 2), overlay(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +HAaAA, world! Nice day. HAAao, world! Nice day. +SAaAA SAAa +SAaAAhing else. Nothing here. SAAathing else. Nothing here. +aAaAA aAAa +aAaAAAaA aAAaAAaA +aAaAAAaA aAAaaAaA +aAaAAAaA aAAaaAaAaaAaaAaAaaAaaAaA +aAaAAba aAAaCbA +bAaAAAbA aAAa +eAaAA eAAa +kAaAAn sAAaing +İAaAA iAAa +İAaAA İAAa + + +-- !query +select format_string(format, utf8_binary, utf8_lcase) from t8 +-- !query schema +struct +-- !query output +abCdEabCdE + + +-- !query +select format_string(format collate utf8_lcase, utf8_lcase, utf8_binary collate utf8_lcase, 3), format_string(format, utf8_lcase collate utf8_binary, utf8_binary) from t8 +-- !query schema +struct +-- !query output +abCdEabCdE abCdEabCdE + + +-- !query +select format_string(format, utf8_binary, utf8_lcase) from t8 +-- !query schema +struct +-- !query output +abCdEabCdE + + +-- !query +select soundex(utf8_binary), soundex(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +A000 A000 +A000 A000 +A000 A000 +A120 A120 +A123 A123 +B110 A000 +E130 E130 +H464 H464 +K350 S352 +S162 S400 +S535 S535 +İo I000 +İo İo + + +-- !query +select soundex(utf8_binary collate utf8_lcase), soundex(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +A000 A000 +A000 A000 +A000 A000 +A120 A120 +A123 A123 +B110 A000 +E130 E130 +H464 H464 +K350 S352 +S162 S400 +S535 S535 +İo I000 +İo İo + + +-- !query +select length(utf8_binary), length(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +2 2 +2 3 +23 23 +29 29 +3 3 +4 4 +5 3 +6 7 +7 7 +8 1 +8 24 +8 8 +8 8 + + +-- !query +select length(utf8_binary collate utf8_lcase), length(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +2 2 +2 3 +23 23 +29 29 +3 3 +4 4 +5 3 +6 7 +7 7 +8 1 +8 24 +8 8 +8 8 + + +-- !query +select bit_length(utf8_binary), bit_length(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +184 184 +232 232 +24 24 +24 24 +24 32 +32 32 +40 24 +48 56 +56 56 +64 192 +64 64 +64 64 +64 8 + + +-- !query +select bit_length(utf8_binary collate utf8_lcase), bit_length(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +184 184 +232 232 +24 24 +24 24 +24 32 +32 32 +40 24 +48 56 +56 56 +64 192 +64 64 +64 64 +64 8 + + +-- !query +select octet_length(utf8_binary), octet_length(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +23 23 +29 29 +3 3 +3 3 +3 4 +4 4 +5 3 +6 7 +7 7 +8 1 +8 24 +8 8 +8 8 + + +-- !query +select octet_length(utf8_binary collate utf8_lcase), octet_length(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +23 23 +29 29 +3 3 +3 3 +3 4 +4 4 +5 3 +6 7 +7 7 +8 1 +8 24 +8 8 +8 8 + + +-- !query +select luhn_check(num) from t9 +-- !query schema +struct +-- !query output +false +false + + +-- !query +select levenshtein(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select levenshtein(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +1 + + +-- !query +select levenshtein(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select levenshtein(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +1 +16 +2 +4 +4 +4 +8 + + +-- !query +select levenshtein(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +1 +16 +2 +4 +4 +4 +8 + + +-- !query +select levenshtein(utf8_binary, 'a'), levenshtein(utf8_lcase, 'a') from t5 +-- !query schema +struct +-- !query output +2 2 +2 2 +2 3 +22 22 +29 29 +4 3 +4 4 +6 6 +6 7 +7 23 +7 7 +7 7 +8 0 + + +-- !query +select levenshtein(utf8_binary, 'AaAA' collate utf8_lcase, 3), levenshtein(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5 +-- !query schema +struct +-- !query output +-1 -1 +-1 -1 +-1 -1 +-1 -1 +-1 -1 +-1 -1 +-1 -1 +-1 2 +-1 3 +-1 3 +-1 3 +-1 4 +3 3 + + +-- !query +select is_valid_utf8(utf8_binary), is_valid_utf8(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true + + +-- !query +select is_valid_utf8(utf8_binary collate utf8_lcase), is_valid_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true +true true + + +-- !query +select make_valid_utf8(utf8_binary), make_valid_utf8(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select make_valid_utf8(utf8_binary collate utf8_lcase), make_valid_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select validate_utf8(utf8_binary), validate_utf8(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select validate_utf8(utf8_binary collate utf8_lcase), validate_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select try_validate_utf8(utf8_binary), try_validate_utf8(utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select try_validate_utf8(utf8_binary collate utf8_lcase), try_validate_utf8(utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaAaaAaaAaAaaAaaAaA +abc abc +abcdcba aBcDCbA +bbAbAAbA a +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select substr(utf8_binary, 2, 2), substr(utf8_lcase, 2, 2) from t5 +-- !query schema +struct +-- !query output +aA aA +aA aA +aA aA +bA +bc Bc +bc bc +el el +fd fd +it it +o o +o ̇o +om om +pa QL + + +-- !query +select substr(utf8_binary collate utf8_lcase, 2, 2), substr(utf8_lcase collate utf8_binary, 2, 2) from t5 +-- !query schema +struct +-- !query output +aA aA +aA aA +aA aA +bA +bc Bc +bc bc +el el +fd fd +it it +o o +o ̇o +om om +pa QL + + +-- !query +select right(utf8_binary, 2), right(utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output +aA aA +aA aA +aA aA +bA a +ba bA +bc bc +d2 d2 +e. e. +en ng +rk QL +y. y. +İo İo +İo ̇o + + +-- !query +select right(utf8_binary collate utf8_lcase, 2), right(utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +aA aA +aA aA +aA aA +bA a +ba bA +bc bc +d2 d2 +e. e. +en ng +rk QL +y. y. +İo İo +İo ̇o + + +-- !query +select left(utf8_binary, '2' collate utf8_lcase), left(utf8_lcase, 2) from t5 +-- !query schema +struct +-- !query output +He He +So So +Sp SQ +aa aa +aa aa +aa aa +ab aB +ab ab +bb a +ef ef +ki si +İo i̇ +İo İo + + +-- !query +select left(utf8_binary collate utf8_lcase, 2), left(utf8_lcase collate utf8_binary, 2) from t5 +-- !query schema +struct +-- !query output +He He +So So +Sp SQ +aa aa +aa aa +aa aa +ab aB +ab ab +bb a +ef ef +ki si +İo i̇ +İo İo + + +-- !query +select rpad(utf8_binary, 8, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select rpad(s, 8, utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, w +Somethin +SparkSpa +aaAaAAaA +aaAaAAaA +aaAaAAaA +abcabcab +abcdcbaa +bbAbaAbA +efd2efd2 +kittenki +İoİoİoİo +İoİoİoİo + + +-- !query +select rpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select rpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, w +Somethin +SparkSQL +aaAaAAaA +aaAaAAaA +aaAaAAaA +abcabcab +abcdcbaa +bbAbAAbA +efd2efd2 +kittensi +İoi̇oi̇o +İoİoİoİo + + +-- !query +select rpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, w +Somethin +SparkSQL +aaAaAAaA +aaAaAAaA +aaAaAAaA +abcabcab +abcdcbaa +bbAbAAbA +efd2efd2 +kittensi +İoi̇oi̇o +İoİoİoİo + + +-- !query +select rpad(utf8_binary, 8, 'a'), rpad(utf8_lcase, 8, 'a') from t5 +-- !query schema +struct +-- !query output +Hello, w Hello, w +Somethin Somethin +Sparkaaa SQLaaaaa +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaA +abcaaaaa abcaaaaa +abcdcbaa aBcDCbAa +bbAbAAbA aaaaaaaa +efd2aaaa efd2aaaa +kittenaa sitTinga +İoaaaaaa i̇oaaaaa +İoaaaaaa İoaaaaaa + + +-- !query +select rpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), rpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, w Hello, w +Somethin Somethin +SparkAaA SQLAAaAA +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaA +abcAaAAA abcAAaAA +abcdcbaA aBcDCbAA +bbAbAAbA aAAaAAaA +efd2AaAA efd2AAaA +kittenAa sitTingA +İoAaAAAa i̇oAAaAA +İoAaAAAa İoAAaAAa + + +-- !query +select lpad(utf8_binary, 8, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select lpad(s, 8, utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, w +Somethin +SpaSpark +aaAaAAaA +aaAaAAaA +aaAaAAaA +aabcdcba +abcababc +bbAbaAbA +efd2efd2 +kikitten +İoİoİoİo +İoİoİoİo + + +-- !query +select lpad(utf8_binary collate utf8_binary, 8, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select lpad(utf8_binary, 8, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +Hello, w +SQLSpark +Somethin +aaAaAAaA +aaAaAAaA +aaAaAAaA +aabcdcba +abcababc +bbAbAAbA +efd2efd2 +i̇oi̇oİo +sikitten +İoİoİoİo + + +-- !query +select lpad(utf8_binary collate utf8_lcase, 8, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, w +SQLSpark +Somethin +aaAaAAaA +aaAaAAaA +aaAaAAaA +aabcdcba +abcababc +bbAbAAbA +efd2efd2 +i̇oi̇oİo +sikitten +İoİoİoİo + + +-- !query +select lpad(utf8_binary, 8, 'a'), lpad(utf8_lcase, 8, 'a') from t5 +-- !query schema +struct +-- !query output +Hello, w Hello, w +Somethin Somethin +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaA +aaaSpark aaaaaSQL +aaaaaabc aaaaaabc +aaaaaaİo aaaaaaİo +aaaaaaİo aaaaai̇o +aaaaefd2 aaaaefd2 +aabcdcba aaBcDCbA +aakitten asitTing +bbAbAAbA aaaaaaaa + + +-- !query +select lpad(utf8_binary, 8, 'AaAA' collate utf8_lcase), lpad(utf8_lcase, 8, 'AAa' collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +AaAAAabc AAaAAabc +AaAAAaİo AAaAAaİo +AaAAAaİo AAaAAi̇o +AaAAefd2 AAaAefd2 +AaASpark AAaAASQL +Aabcdcba AaBcDCbA +Aakitten AsitTing +Hello, w Hello, w +Somethin Somethin +aaAaAAaA aaAaAAaA +aaAaAAaA aaAaaAaA +aaAaAAaA aaAaaAaA +bbAbAAbA AAaAAaAa + + +-- !query +select locate(utf8_binary, utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select locate(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 + + +-- !query +select locate(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select locate(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 +1 + + +-- !query +select locate(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 3) from t5 +-- !query schema +struct +-- !query output +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +3 + + +-- !query +select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 1 + + +-- !query +select locate(utf8_binary, 'AaAA' collate utf8_lcase, 4), locate(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5 +-- !query schema +struct +-- !query output +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 + + +-- !query +select TRIM(utf8_binary, utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -517,7 +3623,15 @@ org.apache.spark.sql.AnalysisException -- !query -select split_part(str collate utf8_binary, delimiter collate utf8_lcase, partNum) from t5 +select TRIM(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + +-- !query +select TRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -526,45 +3640,220 @@ org.apache.spark.sql.AnalysisException "errorClass" : "COLLATION_MISMATCH.EXPLICIT", "sqlState" : "42P21", "messageParameters" : { - "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" } } -- !query -select split_part(str collate utf8_binary, delimiter collate utf8_binary, partNum) from t5 +select TRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 -- !query schema -struct +struct -- !query output -12 + + + + + + + + +BcDCbA +QL +a +i̇ +sitTing -- !query -drop table t5 +select TRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 -- !query schema -struct<> +struct +-- !query output + + + + + + + + + + + +QL +sitTing + + +-- !query +select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5 +-- !query schema +struct -- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAa +aaAaAAa +aaAaAAa +ab +abcdcba D +bbAbAAb +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + +-- !query +select TRIM('ABc' collate utf8_lcase, utf8_binary), TRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + bc +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +d BcDCb +efd2 efd2 +kitten sitTing +İo i̇o +İo İo -- !query -create table t6 (utf8_binary string collate utf8_binary, utf8_lcase string collate utf8_lcase, threshold int) using parquet +select BTRIM(utf8_binary, utf8_lcase) from t5 -- !query schema struct<> -- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.IMPLICIT", + "sqlState" : "42P21" +} + + +-- !query +select BTRIM(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + +a + + -- !query -insert into t6 values('kitten', 'sitting', 2) +select BTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query schema struct<> -- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "COLLATION_MISMATCH.EXPLICIT", + "sqlState" : "42P21", + "messageParameters" : { + "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + } +} + + +-- !query +select BTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + + + + + + +bbAbAAbA +d +kitte +park +İ + + +-- !query +select BTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + +bbAbAAb +kitte +park +İ + + +-- !query +select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5 +-- !query schema +struct +-- !query output +AB +AB +AB B +ABc ABc +ABc ABc +ABc ABc +ABc ABc +ABc ABc +ABc ABc +Bc Bc +Bc Bc +Bc Bc +Bc Bc + +-- !query +select BTRIM('ABc' collate utf8_lcase, utf8_binary), BTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + AA +ABc AAa +ABc AAa +ABc AAa +ABc AAa +ABc AAa +B AA +Bc +Bc +Bc +Bc AAa +c AA -- !query -select levenshtein(utf8_binary, utf8_lcase) from t6 +select LTRIM(utf8_binary, utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -576,7 +3865,15 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase) from t6 +select LTRIM(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + +-- !query +select LTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -585,21 +3882,93 @@ org.apache.spark.sql.AnalysisException "errorClass" : "COLLATION_MISMATCH.EXPLICIT", "sqlState" : "42P21", "messageParameters" : { - "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" } } -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary) from t6 +select LTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 -- !query schema -struct +struct -- !query output -3 + + + + + + + + +BcDCbA +QL +a +i̇o +sitTing + + +-- !query +select LTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + +QL +sitTing + + +-- !query +select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAaA +aaAaAAaA +aaAaAAaA +abc +abcdcba DCbA +bbAbAAbA +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select LTRIM('ABc' collate utf8_lcase, utf8_binary), LTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + bc +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +dcba BcDCbA +efd2 efd2 +kitten sitTing +İo i̇o +İo İo -- !query -select levenshtein(utf8_binary, utf8_lcase, threshold) from t6 +select RTRIM(utf8_binary, utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -611,7 +3980,15 @@ org.apache.spark.sql.AnalysisException -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_lcase, threshold) from t6 +select RTRIM(s, utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + +-- !query +select RTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5 -- !query schema struct<> -- !query output @@ -620,17 +3997,97 @@ org.apache.spark.sql.AnalysisException "errorClass" : "COLLATION_MISMATCH.EXPLICIT", "sqlState" : "42P21", "messageParameters" : { - "explicitTypes" : "`string`, `string collate UTF8_LCASE`" + "explicitTypes" : "`string collate UTF8_LCASE`, `string`" } } -- !query -select levenshtein(utf8_binary collate utf8_binary, utf8_lcase collate utf8_binary, threshold) from t6 +select RTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5 +-- !query schema +struct +-- !query output + + + + + + + + +SQL +a +aBcDCbA +i̇ +sitTing + + +-- !query +select RTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + + + + + + + +SQL +sitTing + + +-- !query +select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5 +-- !query schema +struct +-- !query output +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +aaAaAAa +aaAaAAa +aaAaAAa +ab +abcdcba aBcD +bbAbAAb +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +select RTRIM('ABc' collate utf8_lcase, utf8_binary), RTRIM('AAa' collate utf8_binary, utf8_lcase) from t5 +-- !query schema +struct +-- !query output + + + + + abc +Hello, world! Nice day. Hello, world! Nice day. +Something else. Nothing here. Something else. Nothing here. +Spark SQL +abcd aBcDCb +efd2 efd2 +kitten sitTing +İo i̇o +İo İo + + +-- !query +drop table t5 -- !query schema -struct +struct<> -- !query output --1 + -- !query @@ -639,3 +4096,35 @@ drop table t6 struct<> -- !query output + + +-- !query +drop table t7 +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table t8 +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table t9 +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table t10 +-- !query schema +struct<> +-- !query output + From d3eb99f79e508d62fdb7e9bc595f0240ac021df5 Mon Sep 17 00:00:00 2001 From: Josh Rosen Date: Fri, 13 Sep 2024 15:07:25 -0700 Subject: [PATCH 029/189] [SPARK-49647][TESTS] Change SharedSparkContext so that its SparkConf loads defaults ### What changes were proposed in this pull request? This PR modifies the `SharedSparkContext` text suite mixin trait so that it instantiates SparkConf as `new SparkConf()` instead of `new SparkConf(loadDefaults = false)`. ### Why are the changes needed? Spark's SBT and Maven builds configure certain test default configurations using system properties, including disabling the Spark UI and lowering Derby metastore durability: https://github.com/apache/spark/blob/08a26bb56cfb48f27c68a79be1e15bc4c9e466e0/project/SparkBuild.scala#L1616-L1633 Most test suites pick up defaults set at this layer. However, the `SharedSparkContext` trait was using `new SparkConf(false)` which bypasses the loading of these defaults. As a result, tests which used this trait don't pick up default configurations and instead try to launch the Spark UI, which may lead to test flakiness due to port binding conflicts. This PR proposes to change this to `new SparkConf()` to address this issue. This change brings `SharedSparkContext` to parity with `SharedSparkSession`, which was already running with `loadDefaults = true`. ### Does this PR introduce _any_ user-facing change? This is a test-only change, so it's not expected to impact users of Apache Spark itself. However, it might possibly impact third-party developers who directly depend on Apache Spark's own test JARs. ### How was this patch tested? Running existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48111 from JoshRosen/SPARK-49647-pick-up-sys-prop-defaults-in-SharedSparkContext-test-mixin. Authored-by: Josh Rosen Signed-off-by: Dongjoon Hyun --- .../src/test/scala/org/apache/spark/SharedSparkContext.scala | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/core/src/test/scala/org/apache/spark/SharedSparkContext.scala b/core/src/test/scala/org/apache/spark/SharedSparkContext.scala index 7106a780b3256..22c6280198c9a 100644 --- a/core/src/test/scala/org/apache/spark/SharedSparkContext.scala +++ b/core/src/test/scala/org/apache/spark/SharedSparkContext.scala @@ -27,7 +27,10 @@ trait SharedSparkContext extends BeforeAndAfterAll with BeforeAndAfterEach { sel def sc: SparkContext = _sc - val conf = new SparkConf(false) + // SPARK-49647: use `SparkConf()` instead of `SparkConf(false)` because we want to + // load defaults from system properties and the classpath, including default test + // settings specified in the SBT and Maven build definitions. + val conf: SparkConf = new SparkConf() /** * Initialize the [[SparkContext]]. Generally, this is just called from beforeAll; however, in From df0e34c5a1c30956cb16e8af5569ed72387b6fc3 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Fri, 13 Sep 2024 18:09:48 -0700 Subject: [PATCH 030/189] [SPARK-49648][DOCS] Update `Configuring Ports for Network Security` section with JWS ### What changes were proposed in this pull request? This PR aims to update `Configuring Ports for Network Security` section of `Security` page with new JWS feature. ### Why are the changes needed? In addition to the existing restriction, Spark 4 can take advantage of new JWS feature. This PR informs it more clearly. https://github.com/apache/spark/blob/08a26bb56cfb48f27c68a79be1e15bc4c9e466e0/docs/security.md?plain=1#L811-L814 ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Manual review. Screenshot 2024-09-13 at 15 04 43 Screenshot 2024-09-13 at 15 04 16 ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48112 from dongjoon-hyun/SPARK-49648. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- docs/security.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/security.md b/docs/security.md index a8f4e4ec53897..b97abfeacf240 100644 --- a/docs/security.md +++ b/docs/security.md @@ -55,7 +55,8 @@ To enable authorization, Spark Master should have `spark.master.rest.filters=org.apache.spark.ui.JWSFilter` and `spark.org.apache.spark.ui.JWSFilter.param.secretKey=BASE64URL-ENCODED-KEY` configurations, and client should provide HTTP `Authorization` header which contains JSON Web Token signed by -the shared secret key. +the shared secret key. Please note that this feature requires a Spark distribution built with +`jjwt` profile. ### YARN @@ -813,6 +814,12 @@ They are generally private services, and should only be accessible within the ne organization that deploys Spark. Access to the hosts and ports used by Spark services should be limited to origin hosts that need to access the services. +However, like the REST Submission port, Spark also supports HTTP `Authorization` header +with a cryptographically signed JSON Web Token (JWT) for all UI ports. +To use it, a user needs the Spark distribution built with `jjwt` profile and to configure +`spark.ui.filters=org.apache.spark.ui.JWSFilter` and +`spark.org.apache.spark.ui.JWSFilter.param.secretKey=BASE64URL-ENCODED-KEY`. + Below are the primary ports that Spark uses for its communication and how to configure those ports. From 2250b35be6a24c777d6fa82b1c6a7a10a6854895 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Fri, 13 Sep 2024 20:49:01 -0700 Subject: [PATCH 031/189] [SPARK-49649][DOCS] Make `docs/index.md` up-to-date for 4.0.0 ### What changes were proposed in this pull request? This PR aims to update Spark documentation landing page (`docs/index.md`) for Apache Spark 4.0.0-preview2 release. ### Why are the changes needed? - [SPARK-45314 Drop Scala 2.12 and make Scala 2.13 by default](https://issues.apache.org/jira/browse/SPARK-45314) - #46228 - #47842 - [SPARK-45923 Spark Kubernetes Operator](https://issues.apache.org/jira/browse/SPARK-45923) ### Does this PR introduce _any_ user-facing change? No because this is a documentation-only change. ### How was this patch tested? Manual review. Screenshot 2024-09-13 at 16 01 55 Screenshot 2024-09-13 at 16 02 09 Screenshot 2024-09-13 at 16 02 38 ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48113 from dongjoon-hyun/SPARK-49649. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- docs/index.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/index.md b/docs/index.md index 7e57eddb6da86..fea62865e2160 100644 --- a/docs/index.md +++ b/docs/index.md @@ -34,9 +34,8 @@ source, visit [Building Spark](building-spark.html). Spark runs on both Windows and UNIX-like systems (e.g. Linux, Mac OS), and it should run on any platform that runs a supported version of Java. This should include JVMs on x86_64 and ARM64. It's easy to run locally on one machine --- all you need is to have `java` installed on your system `PATH`, or the `JAVA_HOME` environment variable pointing to a Java installation. -Spark runs on Java 17/21, Scala 2.13, Python 3.8+, and R 3.5+. -When using the Scala API, it is necessary for applications to use the same version of Scala that Spark was compiled for. -For example, when using Scala 2.13, use Spark compiled for 2.13, and compile code/applications for Scala 2.13 as well. +Spark runs on Java 17/21, Scala 2.13, Python 3.9+, and R 3.5+ (Deprecated). +When using the Scala API, it is necessary for applications to use the same version of Scala that Spark was compiled for. Since Spark 4.0.0, it's Scala 2.13. # Running the Examples and Shell @@ -110,7 +109,7 @@ options for deployment: * [Spark Streaming](streaming-programming-guide.html): processing data streams using DStreams (old API) * [MLlib](ml-guide.html): applying machine learning algorithms * [GraphX](graphx-programming-guide.html): processing graphs -* [SparkR](sparkr.html): processing data with Spark in R +* [SparkR (Deprecated)](sparkr.html): processing data with Spark in R * [PySpark](api/python/getting_started/index.html): processing data with Spark in Python * [Spark SQL CLI](sql-distributed-sql-engine-spark-sql-cli.html): processing data with SQL on the command line @@ -128,10 +127,13 @@ options for deployment: * [Cluster Overview](cluster-overview.html): overview of concepts and components when running on a cluster * [Submitting Applications](submitting-applications.html): packaging and deploying applications * Deployment modes: - * [Amazon EC2](https://github.com/amplab/spark-ec2): scripts that let you launch a cluster on EC2 in about 5 minutes * [Standalone Deploy Mode](spark-standalone.html): launch a standalone cluster quickly without a third-party cluster manager * [YARN](running-on-yarn.html): deploy Spark on top of Hadoop NextGen (YARN) - * [Kubernetes](running-on-kubernetes.html): deploy Spark on top of Kubernetes + * [Kubernetes](running-on-kubernetes.html): deploy Spark apps on top of Kubernetes directly + * [Amazon EC2](https://github.com/amplab/spark-ec2): scripts that let you launch a cluster on EC2 in about 5 minutes +* [Spark Kubernetes Operator](https://github.com/apache/spark-kubernetes-operator): + * [SparkApp](https://github.com/apache/spark-kubernetes-operator/blob/main/examples/pyspark-pi.yaml): deploy Spark apps on top of Kubernetes via [operator patterns](https://kubernetes.io/docs/concepts/extend-kubernetes/operator/) + * [SparkCluster](https://github.com/apache/spark-kubernetes-operator/blob/main/examples/cluster-with-template.yaml): deploy Spark clusters on top of Kubernetes via [operator patterns](https://kubernetes.io/docs/concepts/extend-kubernetes/operator/) **Other Documents:** From 017b0ea71e03339336b5d199ecad4f50961e4948 Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Sat, 14 Sep 2024 12:16:35 +0800 Subject: [PATCH 032/189] [SPARK-49556][SQL] Add SQL pipe syntax for the SELECT operator ### What changes were proposed in this pull request? This PR adds SQL pipe syntax support for the SELECT operator. For example: ``` CREATE TABLE t(x INT, y STRING) USING CSV; INSERT INTO t VALUES (0, 'abc'), (1, 'def'); TABLE t |> SELECT x, y 0 abc 1 def TABLE t |> SELECT x, y |> SELECT x + LENGTH(y) AS z 3 4 (SELECT * FROM t UNION ALL SELECT * FROM t) |> SELECT x + LENGTH(y) AS result 3 3 4 4 TABLE t |> SELECT sum(x) AS result Error: aggregate functions are not allowed in the pipe operator |> SELECT clause; please use the |> AGGREGATE clause instead ``` ### Why are the changes needed? The SQL pipe operator syntax will let users compose queries in a more flexible fashion. ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds a few unit test cases, but mostly relies on golden file test coverage. I did this to make sure the answers are correct as this feature is implemented and also so we can look at the analyzer output plans to ensure they look right as well. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48047 from dtenedor/pipe-select. Authored-by: Daniel Tenedorio Signed-off-by: Wenchen Fan --- .../resources/error/error-conditions.json | 6 + .../spark/sql/catalyst/parser/SqlBaseLexer.g4 | 1 + .../sql/catalyst/parser/SqlBaseParser.g4 | 5 + .../sql/catalyst/expressions/PipeSelect.scala | 47 +++ .../sql/catalyst/parser/AstBuilder.scala | 58 +++- .../sql/catalyst/trees/TreePatterns.scala | 1 + .../sql/errors/QueryCompilationErrors.scala | 8 + .../apache/spark/sql/internal/SQLConf.scala | 9 + .../analyzer-results/pipe-operators.sql.out | 318 ++++++++++++++++++ .../sql-tests/inputs/pipe-operators.sql | 102 ++++++ .../sql-tests/results/pipe-operators.sql.out | 308 +++++++++++++++++ .../sql/execution/SparkSqlParserSuite.scala | 19 +- .../ThriftServerQueryTestSuite.scala | 3 +- 13 files changed, 876 insertions(+), 9 deletions(-) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PipeSelect.scala create mode 100644 sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out create mode 100644 sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql create mode 100644 sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 0a9dcd52ea831..a6d8550716b96 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -3754,6 +3754,12 @@ ], "sqlState" : "42K03" }, + "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION" : { + "message" : [ + "Aggregate function is not allowed when using the pipe operator |> SELECT clause; please use the pipe operator |> AGGREGATE clause instead" + ], + "sqlState" : "0A000" + }, "PIVOT_VALUE_DATA_TYPE_MISMATCH" : { "message" : [ "Invalid pivot value '': value data type does not match pivot column data type ." diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index 9ea213f3bf4a6..96a58b99debeb 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -506,6 +506,7 @@ TILDE: '~'; AMPERSAND: '&'; PIPE: '|'; CONCAT_PIPE: '||'; +OPERATOR_PIPE: '|>'; HAT: '^'; COLON: ':'; DOUBLE_COLON: '::'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 73d5cb55295ab..3ea408ca42703 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -613,6 +613,7 @@ queryTerm operator=INTERSECT setQuantifier? right=queryTerm #setOperation | left=queryTerm {!legacy_setops_precedence_enabled}? operator=(UNION | EXCEPT | SETMINUS) setQuantifier? right=queryTerm #setOperation + | left=queryTerm OPERATOR_PIPE operatorPipeRightSide #operatorPipeStatement ; queryPrimary @@ -1471,6 +1472,10 @@ version | stringLit ; +operatorPipeRightSide + : selectClause + ; + // When `SQL_standard_keyword_behavior=true`, there are 2 kinds of keywords in Spark SQL. // - Reserved keywords: // Keywords that are reserved and can't be used as identifiers for table, view, column, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PipeSelect.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PipeSelect.scala new file mode 100644 index 0000000000000..0b5479cc8f0ee --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PipeSelect.scala @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.expressions + +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction +import org.apache.spark.sql.catalyst.trees.TreePattern.{PIPE_OPERATOR_SELECT, RUNTIME_REPLACEABLE, TreePattern} +import org.apache.spark.sql.errors.QueryCompilationErrors + +/** + * Represents a SELECT clause when used with the |> SQL pipe operator. + * We use this to make sure that no aggregate functions exist in the SELECT expressions. + */ +case class PipeSelect(child: Expression) + extends UnaryExpression with RuntimeReplaceable { + final override val nodePatterns: Seq[TreePattern] = Seq(PIPE_OPERATOR_SELECT, RUNTIME_REPLACEABLE) + override def withNewChildInternal(newChild: Expression): Expression = PipeSelect(newChild) + override lazy val replacement: Expression = { + def visit(e: Expression): Unit = e match { + case a: AggregateFunction => + // If we used the pipe operator |> SELECT clause to specify an aggregate function, this is + // invalid; return an error message instructing the user to use the pipe operator + // |> AGGREGATE clause for this purpose instead. + throw QueryCompilationErrors.pipeOperatorSelectContainsAggregateFunction(a) + case _: WindowExpression => + // Window functions are allowed in pipe SELECT operators, so do not traverse into children. + case _ => + e.children.foreach(visit) + } + visit(child) + child + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index 7ad7d60e70c96..edcb417da123b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -469,7 +469,8 @@ class AstBuilder extends DataTypeAstBuilder ctx.aggregationClause, ctx.havingClause, ctx.windowClause, - plan + plan, + isPipeOperatorSelect = false ) } } @@ -1057,7 +1058,8 @@ class AstBuilder extends DataTypeAstBuilder ctx.aggregationClause, ctx.havingClause, ctx.windowClause, - from + from, + isPipeOperatorSelect = false ) } @@ -1144,7 +1146,8 @@ class AstBuilder extends DataTypeAstBuilder aggregationClause, havingClause, windowClause, - isDistinct = false) + isDistinct = false, + isPipeOperatorSelect = false) ScriptTransformation( string(visitStringLit(transformClause.script)), @@ -1165,6 +1168,8 @@ class AstBuilder extends DataTypeAstBuilder * Add a regular (SELECT) query specification to a logical plan. The query specification * is the core of the logical plan, this is where sourcing (FROM clause), projection (SELECT), * aggregation (GROUP BY ... HAVING ...) and filtering (WHERE) takes place. + * If 'isPipeOperatorSelect' is true, wraps each projected expression with a [[PipeSelect]] + * expression for future validation of the expressions during analysis. * * Note that query hints are ignored (both by the parser and the builder). */ @@ -1176,7 +1181,8 @@ class AstBuilder extends DataTypeAstBuilder aggregationClause: AggregationClauseContext, havingClause: HavingClauseContext, windowClause: WindowClauseContext, - relation: LogicalPlan): LogicalPlan = withOrigin(ctx) { + relation: LogicalPlan, + isPipeOperatorSelect: Boolean): LogicalPlan = withOrigin(ctx) { val isDistinct = selectClause.setQuantifier() != null && selectClause.setQuantifier().DISTINCT() != null @@ -1188,7 +1194,8 @@ class AstBuilder extends DataTypeAstBuilder aggregationClause, havingClause, windowClause, - isDistinct) + isDistinct, + isPipeOperatorSelect) // Hint selectClause.hints.asScala.foldRight(plan)(withHints) @@ -1202,7 +1209,8 @@ class AstBuilder extends DataTypeAstBuilder aggregationClause: AggregationClauseContext, havingClause: HavingClauseContext, windowClause: WindowClauseContext, - isDistinct: Boolean): LogicalPlan = { + isDistinct: Boolean, + isPipeOperatorSelect: Boolean): LogicalPlan = { // Add lateral views. val withLateralView = lateralView.asScala.foldLeft(relation)(withGenerate) @@ -1216,7 +1224,20 @@ class AstBuilder extends DataTypeAstBuilder } def createProject() = if (namedExpressions.nonEmpty) { - Project(namedExpressions, withFilter) + val newProjectList: Seq[NamedExpression] = if (isPipeOperatorSelect) { + // If this is a pipe operator |> SELECT clause, add a [[PipeSelect]] expression wrapping + // each alias in the project list, so the analyzer can check invariants later. + namedExpressions.map { + case a: Alias => + a.withNewChildren(Seq(PipeSelect(a.child))) + .asInstanceOf[NamedExpression] + case other => + other + } + } else { + namedExpressions + } + Project(newProjectList, withFilter) } else { withFilter } @@ -5755,6 +5776,29 @@ class AstBuilder extends DataTypeAstBuilder visitSetVariableImpl(ctx.query(), ctx.multipartIdentifierList(), ctx.assignmentList()) } + override def visitOperatorPipeStatement(ctx: OperatorPipeStatementContext): LogicalPlan = { + visitOperatorPipeRightSide(ctx.operatorPipeRightSide(), plan(ctx.left)) + } + + private def visitOperatorPipeRightSide( + ctx: OperatorPipeRightSideContext, left: LogicalPlan): LogicalPlan = { + if (!SQLConf.get.getConf(SQLConf.OPERATOR_PIPE_SYNTAX_ENABLED)) { + operationNotAllowed("Operator pipe SQL syntax using |>", ctx) + } + Option(ctx.selectClause).map { c => + withSelectQuerySpecification( + ctx = ctx, + selectClause = c, + lateralView = new java.util.ArrayList[LateralViewContext](), + whereClause = null, + aggregationClause = null, + havingClause = null, + windowClause = null, + relation = left, + isPipeOperatorSelect = true) + }.get + } + /** * Check plan for any parameters. * If it finds any throws UNSUPPORTED_FEATURE.PARAMETER_MARKER_IN_UNEXPECTED_STATEMENT. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala index cbbfccfcab5e8..826ac52c2b817 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala @@ -72,6 +72,7 @@ object TreePattern extends Enumeration { val NOT: Value = Value val NULL_CHECK: Value = Value val NULL_LITERAL: Value = Value + val PIPE_OPERATOR_SELECT: Value = Value val SERIALIZE_FROM_OBJECT: Value = Value val OR: Value = Value val OUTER_REFERENCE: Value = Value diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index e4c8c76e958f8..f1f8be3d15751 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -4104,4 +4104,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat messageParameters = Map("functionName" -> functionName) ) } + + def pipeOperatorSelectContainsAggregateFunction(expr: Expression): Throwable = { + new AnalysisException( + errorClass = "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION", + messageParameters = Map( + "expr" -> expr.toString), + origin = expr.origin) + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 5853e4b66dcc0..c3a42dfd62a04 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -5012,6 +5012,15 @@ object SQLConf { .stringConf .createWithDefault("versionAsOf") + val OPERATOR_PIPE_SYNTAX_ENABLED = + buildConf("spark.sql.operatorPipeSyntaxEnabled") + .doc("If true, enable operator pipe syntax for Apache Spark SQL. This uses the operator " + + "pipe marker |> to indicate separation between clauses of SQL in a manner that describes " + + "the sequence of steps that the query performs in a composable fashion.") + .version("4.0.0") + .booleanConf + .createWithDefault(Utils.isTesting) + val LEGACY_PERCENTILE_DISC_CALCULATION = buildConf("spark.sql.legacy.percentileDiscCalculation") .internal() .doc("If true, the old bogus percentile_disc calculation is used. The old calculation " + diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out new file mode 100644 index 0000000000000..ab0635fef048b --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out @@ -0,0 +1,318 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +drop table if exists t +-- !query analysis +DropTable true, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t + + +-- !query +create table t(x int, y string) using csv +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`t`, false + + +-- !query +insert into t values (0, 'abc'), (1, 'def') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/t, false, CSV, [path=file:[not included in comparison]/{warehouse_dir}/t], Append, `spark_catalog`.`default`.`t`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/t), [x, y] ++- Project [cast(col1#x as int) AS x#x, cast(col2#x as string) AS y#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +drop table if exists other +-- !query analysis +DropTable true, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.other + + +-- !query +create table other(a int, b int) using json +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`other`, false + + +-- !query +insert into other values (1, 1), (1, 2), (2, 4) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/other, false, JSON, [path=file:[not included in comparison]/{warehouse_dir}/other], Append, `spark_catalog`.`default`.`other`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/other), [a, b] ++- Project [cast(col1#x as int) AS a#x, cast(col2#x as int) AS b#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +drop table if exists st +-- !query analysis +DropTable true, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.st + + +-- !query +create table st(x int, col struct) using parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`st`, false + + +-- !query +insert into st values (1, (2, 3)) +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/st, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/st], Append, `spark_catalog`.`default`.`st`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/st), [x, col] ++- Project [cast(col1#x as int) AS x#x, named_struct(i1, cast(col2#x.col1 as int), i2, cast(col2#x.col2 as int)) AS col#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +table t +|> select 1 as x +-- !query analysis +Project [pipeselect(1) AS x#x] ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select x, y +-- !query analysis +Project [x#x, y#x] ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select x, y +|> select x + length(y) as z +-- !query analysis +Project [pipeselect((x#x + length(y#x))) AS z#x] ++- Project [x#x, y#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +values (0), (1) tab(col) +|> select col * 2 as result +-- !query analysis +Project [pipeselect((col#x * 2)) AS result#x] ++- SubqueryAlias tab + +- LocalRelation [col#x] + + +-- !query +(select * from t union all select * from t) +|> select x + length(y) as result +-- !query analysis +Project [pipeselect((x#x + length(y#x))) AS result#x] ++- Union false, false + :- Project [x#x, y#x] + : +- SubqueryAlias spark_catalog.default.t + : +- Relation spark_catalog.default.t[x#x,y#x] csv + +- Project [x#x, y#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +(table t + |> select x, y + |> select x) +union all +select x from t where x < 1 +-- !query analysis +Union false, false +:- Project [x#x] +: +- Project [x#x, y#x] +: +- SubqueryAlias spark_catalog.default.t +: +- Relation spark_catalog.default.t[x#x,y#x] csv ++- Project [x#x] + +- Filter (x#x < 1) + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +(select col from st) +|> select col.i1 +-- !query analysis +Project [col#x.i1 AS i1#x] ++- Project [col#x] + +- SubqueryAlias spark_catalog.default.st + +- Relation spark_catalog.default.st[x#x,col#x] parquet + + +-- !query +table st +|> select st.col.i1 +-- !query analysis +Project [col#x.i1 AS i1#x] ++- SubqueryAlias spark_catalog.default.st + +- Relation spark_catalog.default.st[x#x,col#x] parquet + + +-- !query +table t +|> select (select a from other where x = a limit 1) as result +-- !query analysis +Project [pipeselect(scalar-subquery#x [x#x]) AS result#x] +: +- GlobalLimit 1 +: +- LocalLimit 1 +: +- Project [a#x] +: +- Filter (outer(x#x) = a#x) +: +- SubqueryAlias spark_catalog.default.other +: +- Relation spark_catalog.default.other[a#x,b#x] json ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +select (values (0) tab(col) |> select col) as result +-- !query analysis +Project [scalar-subquery#x [] AS result#x] +: +- Project [col#x] +: +- SubqueryAlias tab +: +- LocalRelation [col#x] ++- OneRowRelation + + +-- !query +table t +|> select (select any_value(a) from other where x = a limit 1) as result +-- !query analysis +Project [pipeselect(scalar-subquery#x [x#x]) AS result#x] +: +- GlobalLimit 1 +: +- LocalLimit 1 +: +- Aggregate [any_value(a#x, false) AS any_value(a)#x] +: +- Filter (outer(x#x) = a#x) +: +- SubqueryAlias spark_catalog.default.other +: +- Relation spark_catalog.default.other[a#x,b#x] json ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select x + length(x) as z, z + 1 as plus_one +-- !query analysis +Project [z#x, pipeselect((z#x + 1)) AS plus_one#x] ++- Project [x#x, y#x, pipeselect((x#x + length(cast(x#x as string)))) AS z#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select first_value(x) over (partition by y) as result +-- !query analysis +Project [result#x] ++- Project [x#x, y#x, _we0#x, pipeselect(_we0#x) AS result#x] + +- Window [first_value(x#x, false) windowspecdefinition(y#x, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS _we0#x], [y#x] + +- Project [x#x, y#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +select 1 x, 2 y, 3 z +|> select 1 + sum(x) over (), + avg(y) over (), + x, + avg(x+1) over (partition by y order by z) AS a2 +|> select a2 +-- !query analysis +Project [a2#x] ++- Project [(1 + sum(x) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING))#xL, avg(y) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#x, x#x, a2#x] + +- Project [x#x, y#x, _w1#x, z#x, _we0#xL, avg(y) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#x, _we2#x, (cast(1 as bigint) + _we0#xL) AS (1 + sum(x) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING))#xL, avg(y) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#x, pipeselect(_we2#x) AS a2#x] + +- Window [avg(_w1#x) windowspecdefinition(y#x, z#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS _we2#x], [y#x], [z#x ASC NULLS FIRST] + +- Window [sum(x#x) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS _we0#xL, avg(y#x) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS avg(y) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#x] + +- Project [x#x, y#x, (x#x + 1) AS _w1#x, z#x] + +- Project [1 AS x#x, 2 AS y#x, 3 AS z#x] + +- OneRowRelation + + +-- !query +table t +|> select x, count(*) over () +|> select x +-- !query analysis +Project [x#x] ++- Project [x#x, count(1) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#xL] + +- Project [x#x, count(1) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#xL, count(1) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#xL] + +- Window [count(1) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS count(1) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)#xL] + +- Project [x#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select distinct x, y +-- !query analysis +Distinct ++- Project [x#x, y#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select sum(x) as result +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION", + "sqlState" : "0A000", + "messageParameters" : { + "expr" : "sum(x#x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "sum(x)" + } ] +} + + +-- !query +table t +|> select y, length(y) + sum(x) as result +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION", + "sqlState" : "0A000", + "messageParameters" : { + "expr" : "sum(x#x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 34, + "stopIndex" : 39, + "fragment" : "sum(x)" + } ] +} + + +-- !query +drop table t +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.t + + +-- !query +drop table other +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.other + + +-- !query +drop table st +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.st diff --git a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql new file mode 100644 index 0000000000000..7d0966e7f2095 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql @@ -0,0 +1,102 @@ +-- Prepare some test data. +-------------------------- +drop table if exists t; +create table t(x int, y string) using csv; +insert into t values (0, 'abc'), (1, 'def'); + +drop table if exists other; +create table other(a int, b int) using json; +insert into other values (1, 1), (1, 2), (2, 4); + +drop table if exists st; +create table st(x int, col struct) using parquet; +insert into st values (1, (2, 3)); + +-- Selection operators: positive tests. +--------------------------------------- + +-- Selecting a constant. +table t +|> select 1 as x; + +-- Selecting attributes. +table t +|> select x, y; + +-- Chained pipe SELECT operators. +table t +|> select x, y +|> select x + length(y) as z; + +-- Using the VALUES list as the source relation. +values (0), (1) tab(col) +|> select col * 2 as result; + +-- Using a table subquery as the source relation. +(select * from t union all select * from t) +|> select x + length(y) as result; + +-- Enclosing the result of a pipe SELECT operation in a table subquery. +(table t + |> select x, y + |> select x) +union all +select x from t where x < 1; + +-- Selecting struct fields. +(select col from st) +|> select col.i1; + +table st +|> select st.col.i1; + +-- Expression subqueries in the pipe operator SELECT list. +table t +|> select (select a from other where x = a limit 1) as result; + +-- Pipe operator SELECT inside expression subqueries. +select (values (0) tab(col) |> select col) as result; + +-- Aggregations are allowed within expression subqueries in the pipe operator SELECT list as long as +-- no aggregate functions exist in the top-level select list. +table t +|> select (select any_value(a) from other where x = a limit 1) as result; + +-- Lateral column aliases in the pipe operator SELECT list. +table t +|> select x + length(x) as z, z + 1 as plus_one; + +-- Window functions are allowed in the pipe operator SELECT list. +table t +|> select first_value(x) over (partition by y) as result; + +select 1 x, 2 y, 3 z +|> select 1 + sum(x) over (), + avg(y) over (), + x, + avg(x+1) over (partition by y order by z) AS a2 +|> select a2; + +table t +|> select x, count(*) over () +|> select x; + +-- DISTINCT is supported. +table t +|> select distinct x, y; + +-- Selection operators: negative tests. +--------------------------------------- + +-- Aggregate functions are not allowed in the pipe operator SELECT list. +table t +|> select sum(x) as result; + +table t +|> select y, length(y) + sum(x) as result; + +-- Cleanup. +----------- +drop table t; +drop table other; +drop table st; diff --git a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out new file mode 100644 index 0000000000000..7e0b7912105c2 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out @@ -0,0 +1,308 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +drop table if exists t +-- !query schema +struct<> +-- !query output + + + +-- !query +create table t(x int, y string) using csv +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into t values (0, 'abc'), (1, 'def') +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table if exists other +-- !query schema +struct<> +-- !query output + + + +-- !query +create table other(a int, b int) using json +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into other values (1, 1), (1, 2), (2, 4) +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table if exists st +-- !query schema +struct<> +-- !query output + + + +-- !query +create table st(x int, col struct) using parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +insert into st values (1, (2, 3)) +-- !query schema +struct<> +-- !query output + + + +-- !query +table t +|> select 1 as x +-- !query schema +struct +-- !query output +1 +1 + + +-- !query +table t +|> select x, y +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> select x, y +|> select x + length(y) as z +-- !query schema +struct +-- !query output +3 +4 + + +-- !query +values (0), (1) tab(col) +|> select col * 2 as result +-- !query schema +struct +-- !query output +0 +2 + + +-- !query +(select * from t union all select * from t) +|> select x + length(y) as result +-- !query schema +struct +-- !query output +3 +3 +4 +4 + + +-- !query +(table t + |> select x, y + |> select x) +union all +select x from t where x < 1 +-- !query schema +struct +-- !query output +0 +0 +1 + + +-- !query +(select col from st) +|> select col.i1 +-- !query schema +struct +-- !query output +2 + + +-- !query +table st +|> select st.col.i1 +-- !query schema +struct +-- !query output +2 + + +-- !query +table t +|> select (select a from other where x = a limit 1) as result +-- !query schema +struct +-- !query output +1 +NULL + + +-- !query +select (values (0) tab(col) |> select col) as result +-- !query schema +struct +-- !query output +0 + + +-- !query +table t +|> select (select any_value(a) from other where x = a limit 1) as result +-- !query schema +struct +-- !query output +1 +NULL + + +-- !query +table t +|> select x + length(x) as z, z + 1 as plus_one +-- !query schema +struct +-- !query output +1 2 +2 3 + + +-- !query +table t +|> select first_value(x) over (partition by y) as result +-- !query schema +struct +-- !query output +0 +1 + + +-- !query +select 1 x, 2 y, 3 z +|> select 1 + sum(x) over (), + avg(y) over (), + x, + avg(x+1) over (partition by y order by z) AS a2 +|> select a2 +-- !query schema +struct +-- !query output +2.0 + + +-- !query +table t +|> select x, count(*) over () +|> select x +-- !query schema +struct +-- !query output +0 +1 + + +-- !query +table t +|> select distinct x, y +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> select sum(x) as result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION", + "sqlState" : "0A000", + "messageParameters" : { + "expr" : "sum(x#x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "sum(x)" + } ] +} + + +-- !query +table t +|> select y, length(y) + sum(x) as result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "PIPE_OPERATOR_SELECT_CONTAINS_AGGREGATE_FUNCTION", + "sqlState" : "0A000", + "messageParameters" : { + "expr" : "sum(x#x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 34, + "stopIndex" : 39, + "fragment" : "sum(x)" + } ] +} + + +-- !query +drop table t +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table other +-- !query schema +struct<> +-- !query output + + + +-- !query +drop table st +-- !query schema +struct<> +-- !query output + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala index decfb5555dd87..a80444feb68ae 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala @@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.analysis.{AnalysisTest, UnresolvedAlias, Un import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Concat, GreaterThan, Literal, NullsFirst, SortOrder, UnresolvedWindowExpression, UnspecifiedFrame, WindowSpecDefinition, WindowSpecReference} import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.trees.TreePattern.{LOCAL_RELATION, PROJECT, UNRESOLVED_RELATION} import org.apache.spark.sql.connector.catalog.TableCatalog import org.apache.spark.sql.execution.command._ import org.apache.spark.sql.execution.datasources.{CreateTempViewUsing, RefreshResource} -import org.apache.spark.sql.internal.StaticSQLConf +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StringType import org.apache.spark.util.ArrayImplicits._ @@ -880,4 +881,20 @@ class SparkSqlParserSuite extends AnalysisTest with SharedSparkSession { parser.parsePlan("SELECT\u30001") // Unicode ideographic space } // scalastyle:on + + test("Operator pipe SQL syntax") { + withSQLConf(SQLConf.OPERATOR_PIPE_SYNTAX_ENABLED.key -> "true") { + // Basic selection. + // Here we check that every parsed plan contains a projection and a source relation or + // inline table. + def checkPipeSelect(query: String): Unit = { + val plan: LogicalPlan = parser.parsePlan(query) + assert(plan.containsPattern(PROJECT)) + assert(plan.containsAnyPattern(UNRESOLVED_RELATION, LOCAL_RELATION)) + } + checkPipeSelect("TABLE t |> SELECT 1 AS X") + checkPipeSelect("TABLE t |> SELECT 1 AS X, 2 AS Y |> SELECT X + Y AS Z") + checkPipeSelect("VALUES (0), (1) tab(col) |> SELECT col * 2 AS result") + } + } } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala index 026b2388c593c..331572e62f566 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala @@ -103,7 +103,8 @@ class ThriftServerQueryTestSuite extends SQLQueryTestSuite with SharedThriftServ // SPARK-42921 "timestampNTZ/datetime-special-ansi.sql", // SPARK-47264 - "collations.sql" + "collations.sql", + "pipe-operators.sql" ) override def runQueries( From fa6a0786bb4b23a895e68a721df9ee88684c4fab Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Sat, 14 Sep 2024 17:57:35 -0700 Subject: [PATCH 033/189] Revert "[SPARK-49531][PYTHON][CONNECT] Support line plot with plotly backend" This reverts commit 3b8dddac65bce6f88f51e23e777d521d65fa3373. --- dev/sparktestsupport/modules.py | 4 - python/pyspark/errors/error-conditions.json | 5 - python/pyspark/sql/classic/dataframe.py | 5 - python/pyspark/sql/connect/dataframe.py | 5 - python/pyspark/sql/dataframe.py | 27 ---- python/pyspark/sql/plot/__init__.py | 21 --- python/pyspark/sql/plot/core.py | 135 ------------------ python/pyspark/sql/plot/plotly.py | 30 ---- .../tests/connect/test_parity_frame_plot.py | 36 ----- .../connect/test_parity_frame_plot_plotly.py | 36 ----- python/pyspark/sql/tests/plot/__init__.py | 16 --- .../pyspark/sql/tests/plot/test_frame_plot.py | 79 ---------- .../sql/tests/plot/test_frame_plot_plotly.py | 64 --------- python/pyspark/sql/utils.py | 17 --- python/pyspark/testing/sqlutils.py | 7 - .../apache/spark/sql/internal/SQLConf.scala | 27 ---- 16 files changed, 514 deletions(-) delete mode 100644 python/pyspark/sql/plot/__init__.py delete mode 100644 python/pyspark/sql/plot/core.py delete mode 100644 python/pyspark/sql/plot/plotly.py delete mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot.py delete mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py delete mode 100644 python/pyspark/sql/tests/plot/__init__.py delete mode 100644 python/pyspark/sql/tests/plot/test_frame_plot.py delete mode 100644 python/pyspark/sql/tests/plot/test_frame_plot_plotly.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index b9a4bed715f67..34fbb8450d544 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -548,8 +548,6 @@ def __hash__(self): "pyspark.sql.tests.test_udtf", "pyspark.sql.tests.test_utils", "pyspark.sql.tests.test_resources", - "pyspark.sql.tests.plot.test_frame_plot", - "pyspark.sql.tests.plot.test_frame_plot_plotly", ], ) @@ -1053,8 +1051,6 @@ def __hash__(self): "pyspark.sql.tests.connect.test_parity_arrow_cogrouped_map", "pyspark.sql.tests.connect.test_parity_python_datasource", "pyspark.sql.tests.connect.test_parity_python_streaming_datasource", - "pyspark.sql.tests.connect.test_parity_frame_plot", - "pyspark.sql.tests.connect.test_parity_frame_plot_plotly", "pyspark.sql.tests.connect.test_utils", "pyspark.sql.tests.connect.client.test_artifact", "pyspark.sql.tests.connect.client.test_artifact_localcluster", diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 92aeb15e21d1b..4061d024a83cd 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -1088,11 +1088,6 @@ "Function `` should use only POSITIONAL or POSITIONAL OR KEYWORD arguments." ] }, - "UNSUPPORTED_PLOT_BACKEND": { - "message": [ - "`` is not supported, it should be one of the values from " - ] - }, "UNSUPPORTED_SIGNATURE": { "message": [ "Unsupported signature: ." diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index d174f7774cc57..91b9591625904 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -58,7 +58,6 @@ from pyspark.sql.classic.column import _to_seq, _to_list, _to_java_column from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter -from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.types import ( StructType, @@ -1863,10 +1862,6 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: messageParameters={"member": "queryExecution"}, ) - @property - def plot(self) -> PySparkPlotAccessor: - return PySparkPlotAccessor(self) - class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index e3b1d35b2d5d6..768abd655d497 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -83,7 +83,6 @@ UnresolvedStar, ) from pyspark.sql.connect.functions import builtin as F -from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema from pyspark.sql.pandas.functions import _validate_pandas_udf # type: ignore[attr-defined] @@ -2240,10 +2239,6 @@ def rdd(self) -> "RDD[Row]": def executionInfo(self) -> Optional["ExecutionInfo"]: return self._execution_info - @property - def plot(self) -> PySparkPlotAccessor: - return PySparkPlotAccessor(self) - class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index 7748510258eaa..ef35b73332572 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -39,7 +39,6 @@ from pyspark.sql.column import Column from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter -from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.types import StructType, Row from pyspark.sql.utils import dispatch_df_method @@ -6395,32 +6394,6 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: """ ... - @property - def plot(self) -> PySparkPlotAccessor: - """ - Returns a :class:`PySparkPlotAccessor` for plotting functions. - - .. versionadded:: 4.0.0 - - Returns - ------- - :class:`PySparkPlotAccessor` - - Notes - ----- - This API is experimental. - - Examples - -------- - >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] - >>> columns = ["category", "int_val", "float_val"] - >>> df = spark.createDataFrame(data, columns) - >>> type(df.plot) - - >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP - """ - ... - class DataFrameNaFunctions: """Functionality for working with missing data in :class:`DataFrame`. diff --git a/python/pyspark/sql/plot/__init__.py b/python/pyspark/sql/plot/__init__.py deleted file mode 100644 index 6da07061b2a09..0000000000000 --- a/python/pyspark/sql/plot/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -""" -This package includes the plotting APIs for PySpark DataFrame. -""" -from pyspark.sql.plot.core import * # noqa: F403, F401 diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py deleted file mode 100644 index baee610dc6bd0..0000000000000 --- a/python/pyspark/sql/plot/core.py +++ /dev/null @@ -1,135 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from typing import Any, TYPE_CHECKING, Optional, Union -from types import ModuleType -from pyspark.errors import PySparkRuntimeError, PySparkValueError -from pyspark.sql.utils import require_minimum_plotly_version - - -if TYPE_CHECKING: - from pyspark.sql import DataFrame - import pandas as pd - from plotly.graph_objs import Figure - - -class PySparkTopNPlotBase: - def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": - from pyspark.sql import SparkSession - - session = SparkSession.getActiveSession() - if session is None: - raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) - - max_rows = int( - session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] - ) - pdf = sdf.limit(max_rows + 1).toPandas() - - self.partial = False - if len(pdf) > max_rows: - self.partial = True - pdf = pdf.iloc[:max_rows] - - return pdf - - -class PySparkSampledPlotBase: - def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": - from pyspark.sql import SparkSession - - session = SparkSession.getActiveSession() - if session is None: - raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) - - sample_ratio = session.conf.get("spark.sql.pyspark.plotting.sample_ratio") - max_rows = int( - session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] - ) - - if sample_ratio is None: - fraction = 1 / (sdf.count() / max_rows) - fraction = min(1.0, fraction) - else: - fraction = float(sample_ratio) - - sampled_sdf = sdf.sample(fraction=fraction) - pdf = sampled_sdf.toPandas() - - return pdf - - -class PySparkPlotAccessor: - plot_data_map = { - "line": PySparkSampledPlotBase().get_sampled, - } - _backends = {} # type: ignore[var-annotated] - - def __init__(self, data: "DataFrame"): - self.data = data - - def __call__( - self, kind: str = "line", backend: Optional[str] = None, **kwargs: Any - ) -> "Figure": - plot_backend = PySparkPlotAccessor._get_plot_backend(backend) - - return plot_backend.plot_pyspark(self.data, kind=kind, **kwargs) - - @staticmethod - def _get_plot_backend(backend: Optional[str] = None) -> ModuleType: - backend = backend or "plotly" - - if backend in PySparkPlotAccessor._backends: - return PySparkPlotAccessor._backends[backend] - - if backend == "plotly": - require_minimum_plotly_version() - else: - raise PySparkValueError( - errorClass="UNSUPPORTED_PLOT_BACKEND", - messageParameters={"backend": backend, "supported_backends": ", ".join(["plotly"])}, - ) - from pyspark.sql.plot import plotly as module - - return module - - def line(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": - """ - Plot DataFrame as lines. - - Parameters - ---------- - x : str - Name of column to use for the horizontal axis. - y : str or list of str - Name(s) of the column(s) to use for the vertical axis. Multiple columns can be plotted. - **kwds : optional - Additional keyword arguments. - - Returns - ------- - :class:`plotly.graph_objs.Figure` - - Examples - -------- - >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] - >>> columns = ["category", "int_val", "float_val"] - >>> df = spark.createDataFrame(data, columns) - >>> df.plot.line(x="category", y="int_val") # doctest: +SKIP - >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP - """ - return self(kind="line", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/plot/plotly.py b/python/pyspark/sql/plot/plotly.py deleted file mode 100644 index 5efc19476057f..0000000000000 --- a/python/pyspark/sql/plot/plotly.py +++ /dev/null @@ -1,30 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from typing import TYPE_CHECKING, Any - -from pyspark.sql.plot import PySparkPlotAccessor - -if TYPE_CHECKING: - from pyspark.sql import DataFrame - from plotly.graph_objs import Figure - - -def plot_pyspark(data: "DataFrame", kind: str, **kwargs: Any) -> "Figure": - import plotly - - return plotly.plot(PySparkPlotAccessor.plot_data_map[kind](data), kind, **kwargs) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py deleted file mode 100644 index c69e438bf7eb0..0000000000000 --- a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py +++ /dev/null @@ -1,36 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from pyspark.testing.connectutils import ReusedConnectTestCase -from pyspark.sql.tests.plot.test_frame_plot import DataFramePlotTestsMixin - - -class FramePlotParityTests(DataFramePlotTestsMixin, ReusedConnectTestCase): - pass - - -if __name__ == "__main__": - import unittest - from pyspark.sql.tests.connect.test_parity_frame_plot import * # noqa: F401 - - try: - import xmlrunner # type: ignore[import] - - testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) - except ImportError: - testRunner = None - unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py deleted file mode 100644 index 78508fe533379..0000000000000 --- a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py +++ /dev/null @@ -1,36 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from pyspark.testing.connectutils import ReusedConnectTestCase -from pyspark.sql.tests.plot.test_frame_plot_plotly import DataFramePlotPlotlyTestsMixin - - -class FramePlotPlotlyParityTests(DataFramePlotPlotlyTestsMixin, ReusedConnectTestCase): - pass - - -if __name__ == "__main__": - import unittest - from pyspark.sql.tests.connect.test_parity_frame_plot_plotly import * # noqa: F401 - - try: - import xmlrunner # type: ignore[import] - - testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) - except ImportError: - testRunner = None - unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/__init__.py b/python/pyspark/sql/tests/plot/__init__.py deleted file mode 100644 index cce3acad34a49..0000000000000 --- a/python/pyspark/sql/tests/plot/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# diff --git a/python/pyspark/sql/tests/plot/test_frame_plot.py b/python/pyspark/sql/tests/plot/test_frame_plot.py deleted file mode 100644 index 19ef53e46b2f4..0000000000000 --- a/python/pyspark/sql/tests/plot/test_frame_plot.py +++ /dev/null @@ -1,79 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from pyspark.errors import PySparkValueError -from pyspark.sql import Row -from pyspark.sql.plot import PySparkSampledPlotBase, PySparkTopNPlotBase -from pyspark.testing.sqlutils import ReusedSQLTestCase - - -class DataFramePlotTestsMixin: - def test_backend(self): - accessor = self.spark.range(2).plot - backend = accessor._get_plot_backend() - self.assertEqual(backend.__name__, "pyspark.sql.plot.plotly") - - with self.assertRaises(PySparkValueError) as pe: - accessor._get_plot_backend("matplotlib") - - self.check_error( - exception=pe.exception, - errorClass="UNSUPPORTED_PLOT_BACKEND", - messageParameters={"backend": "matplotlib", "supported_backends": "plotly"}, - ) - - def test_topn_max_rows(self): - try: - self.spark.conf.set("spark.sql.pyspark.plotting.max_rows", "1000") - sdf = self.spark.range(2500) - pdf = PySparkTopNPlotBase().get_top_n(sdf) - self.assertEqual(len(pdf), 1000) - finally: - self.spark.conf.unset("spark.sql.pyspark.plotting.max_rows") - - def test_sampled_plot_with_ratio(self): - try: - self.spark.conf.set("spark.sql.pyspark.plotting.sample_ratio", "0.5") - data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2500)] - sdf = self.spark.createDataFrame(data) - pdf = PySparkSampledPlotBase().get_sampled(sdf) - self.assertEqual(round(len(pdf) / 2500, 1), 0.5) - finally: - self.spark.conf.unset("spark.sql.pyspark.plotting.sample_ratio") - - def test_sampled_plot_with_max_rows(self): - data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2000)] - sdf = self.spark.createDataFrame(data) - pdf = PySparkSampledPlotBase().get_sampled(sdf) - self.assertEqual(round(len(pdf) / 2000, 1), 0.5) - - -class DataFramePlotTests(DataFramePlotTestsMixin, ReusedSQLTestCase): - pass - - -if __name__ == "__main__": - import unittest - from pyspark.sql.tests.plot.test_frame_plot import * # noqa: F401 - - try: - import xmlrunner - - testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) - except ImportError: - testRunner = None - unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py deleted file mode 100644 index 72a3ed267d192..0000000000000 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ /dev/null @@ -1,64 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import unittest -import pyspark.sql.plot # noqa: F401 -from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message - - -@unittest.skipIf(not have_plotly, plotly_requirement_message) -class DataFramePlotPlotlyTestsMixin: - @property - def sdf(self): - data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] - columns = ["category", "int_val", "float_val"] - return self.spark.createDataFrame(data, columns) - - def _check_fig_data(self, fig_data, expected_x, expected_y, expected_name=""): - self.assertEqual(fig_data["mode"], "lines") - self.assertEqual(fig_data["type"], "scatter") - self.assertEqual(fig_data["xaxis"], "x") - self.assertEqual(list(fig_data["x"]), expected_x) - self.assertEqual(fig_data["yaxis"], "y") - self.assertEqual(list(fig_data["y"]), expected_y) - self.assertEqual(fig_data["name"], expected_name) - - def test_line_plot(self): - # single column as vertical axis - fig = self.sdf.plot(kind="line", x="category", y="int_val") - self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20]) - - # multiple columns as vertical axis - fig = self.sdf.plot.line(x="category", y=["int_val", "float_val"]) - self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") - self._check_fig_data(fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") - - -class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): - pass - - -if __name__ == "__main__": - from pyspark.sql.tests.plot.test_frame_plot_plotly import * # noqa: F401 - - try: - import xmlrunner - - testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) - except ImportError: - testRunner = None - unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py index 5d9ec92cbc830..11b91612419a3 100644 --- a/python/pyspark/sql/utils.py +++ b/python/pyspark/sql/utils.py @@ -41,7 +41,6 @@ PythonException, UnknownException, SparkUpgradeException, - PySparkImportError, PySparkNotImplementedError, PySparkRuntimeError, ) @@ -116,22 +115,6 @@ def require_test_compiled() -> None: ) -def require_minimum_plotly_version() -> None: - """Raise ImportError if plotly is not installed""" - minimum_plotly_version = "4.8" - - try: - import plotly # noqa: F401 - except ImportError as error: - raise PySparkImportError( - errorClass="PACKAGE_NOT_INSTALLED", - messageParameters={ - "package_name": "plotly", - "minimum_version": str(minimum_plotly_version), - }, - ) from error - - class ForeachBatchFunction: """ This is the Python implementation of Java interface 'ForeachBatchFunction'. This wraps diff --git a/python/pyspark/testing/sqlutils.py b/python/pyspark/testing/sqlutils.py index 00ad40e68bd7c..9f07c44c084cf 100644 --- a/python/pyspark/testing/sqlutils.py +++ b/python/pyspark/testing/sqlutils.py @@ -48,13 +48,6 @@ except Exception as e: test_not_compiled_message = str(e) -plotly_requirement_message = None -try: - import plotly -except ImportError as e: - plotly_requirement_message = str(e) -have_plotly = plotly_requirement_message is None - from pyspark.sql import SparkSession from pyspark.sql.types import ArrayType, DoubleType, UserDefinedType, Row from pyspark.testing.utils import ReusedPySparkTestCase, PySparkErrorTestUtils diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index c3a42dfd62a04..094fb8f050bc8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3169,29 +3169,6 @@ object SQLConf { .version("4.0.0") .fallbackConf(Python.PYTHON_WORKER_FAULTHANLDER_ENABLED) - val PYSPARK_PLOT_MAX_ROWS = - buildConf("spark.sql.pyspark.plotting.max_rows") - .doc( - "The visual limit on top-n-based plots. If set to 1000, the first 1000 data points " + - "will be used for plotting.") - .version("4.0.0") - .intConf - .createWithDefault(1000) - - val PYSPARK_PLOT_SAMPLE_RATIO = - buildConf("spark.sql.pyspark.plotting.sample_ratio") - .doc( - "The proportion of data that will be plotted for sample-based plots. It is determined " + - "based on spark.sql.pyspark.plotting.max_rows if not explicitly set." - ) - .version("4.0.0") - .doubleConf - .checkValue( - ratio => ratio >= 0.0 && ratio <= 1.0, - "The value should be between 0.0 and 1.0 inclusive." - ) - .createOptional - val ARROW_SPARKR_EXECUTION_ENABLED = buildConf("spark.sql.execution.arrow.sparkr.enabled") .doc("When true, make use of Apache Arrow for columnar data transfers in SparkR. " + @@ -5887,10 +5864,6 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pythonUDFWorkerFaulthandlerEnabled: Boolean = getConf(PYTHON_UDF_WORKER_FAULTHANLDER_ENABLED) - def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS) - - def pysparkPlotSampleRatio: Option[Double] = getConf(PYSPARK_PLOT_SAMPLE_RATIO) - def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED) def arrowPySparkFallbackEnabled: Boolean = getConf(ARROW_PYSPARK_FALLBACK_ENABLED) From 1346531ccc6ee814d5b357158a4c4aed2bf1d573 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Sat, 14 Sep 2024 21:10:29 -0700 Subject: [PATCH 034/189] [SPARK-48355][SQL][TESTS][FOLLOWUP] Disable a test case failing on non-ANSI mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? This PR is a follow-up of https://github.com/apache/spark/pull/47672 to disable a test case failing on non-ANSI mode. ### Why are the changes needed? To recover non-ANSI CI. - https://github.com/apache/spark/actions/workflows/build_non_ansi.yml ### Does this PR introduce _any_ user-facing change? No, this is a test-only change. ### How was this patch tested? Manual review. ``` $ SPARK_ANSI_SQL_MODE=false build/sbt "sql/testOnly *.SqlScriptingInterpreterSuite" ... [info] - simple case mismatched types !!! IGNORED !!! [info] All tests passed. [success] Total time: 24 s, completed Sep 14, 2024, 7:51:15 PM ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48115 from dongjoon-hyun/SPARK-48355. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../spark/sql/scripting/SqlScriptingInterpreterSuite.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala index 3fad99eba509a..bc2adec5be3d5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala @@ -701,7 +701,8 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { verifySqlScriptResult(commands, expected) } - test("simple case mismatched types") { + // This is disabled because it fails in non-ANSI mode + ignore("simple case mismatched types") { val commands = """ |BEGIN From 931ab065df3952487028316ebd49c2895d947bf2 Mon Sep 17 00:00:00 2001 From: "zhipeng.mao" Date: Sun, 15 Sep 2024 13:35:00 +0800 Subject: [PATCH 035/189] [SPARK-48824][SQL] Add Identity Column SQL syntax ### What changes were proposed in this pull request? Add SQL support for creating identity columns. Users can specify a column `GENERATED ALWAYS AS IDENTITY(identityColumnSpec)` , where identity values are **always** generated by the system, or `GENERATED BY DEFAULT AS IDENTITY(identityColumnSpec)`, where users can specify the identity values. Users can optionally specify the starting value of the column (default = 1) and the increment/step of the column (default = 1). Also we allow both `START WITH INCREMENT BY ` and `INCREMENT BY START WITH ` It allows flexible ordering of the increment and starting values, as both variants are used in the wild by other systems (e.g. [PostgreSQL](https://www.postgresql.org/docs/current/sql-createsequence.html) [Oracle](https://docs.oracle.com/en/database/oracle/oracle-database/23/sqlrf/CREATE-SEQUENCE.html#GUID-E9C78A8C-615A-4757-B2A8-5E6EFB130571)). For example, we can define ``` CREATE TABLE default.example ( id LONG GENERATED ALWAYS AS IDENTITY, id1 LONG GENERATED ALWAYS AS IDENTITY(), id2 LONG GENERATED BY DEFAULT AS IDENTITY(START WITH 0), id3 LONG GENERATED ALWAYS AS IDENTITY(INCREMENT BY 2), id4 LONG GENERATED BY DEFAULT AS IDENTITY(START WITH 0 INCREMENT BY -10), id5 LONG GENERATED ALWAYS AS IDENTITY(INCREMENT BY 2 START WITH -8), value LONG ) ``` This will enable defining identity columns in Spark SQL for data sources that support it. To be more specific this PR - Adds parser support for GENERATED { BY DEFAULT | ALWAYS } AS IDENTITY in create/replace table statements. Identity column specifications are temporarily stored in the field's metadata, and then are parsed/verified in DataSourceV2Strategy and used to instantiate v2 [Column] - Adds TableCatalog::capabilities() and TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS This will be used to determine whether to allow specifying identity columns or whether to throw an exception. ### Why are the changes needed? A SQL API is needed to create Identity Columns. ### Does this PR introduce _any_ user-facing change? It allows the aforementioned SQL syntax to create identity columns in a table. ### How was this patch tested? Positive and negative unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47614 from zhipengmao-db/zhipengmao-db/SPARK-48824-id-syntax. Authored-by: zhipeng.mao Signed-off-by: Wenchen Fan --- .../resources/error/error-conditions.json | 24 ++ docs/sql-ref-ansi-compliance.md | 2 + .../spark/sql/catalyst/parser/SqlBaseLexer.g4 | 2 + .../sql/catalyst/parser/SqlBaseParser.g4 | 21 +- .../connector/catalog/IdentityColumnSpec.java | 88 ++++++++ .../spark/sql/errors/QueryParsingErrors.scala | 19 ++ .../spark/sql/connector/catalog/Column.java | 24 +- .../catalog/TableCatalogCapability.java | 20 +- .../sql/catalyst/parser/AstBuilder.scala | 66 +++++- .../plans/logical/ColumnDefinition.scala | 68 ++++-- .../sql/catalyst/util/IdentityColumn.scala | 78 +++++++ .../sql/connector/catalog/CatalogV2Util.scala | 47 +++- .../sql/internal/connector/ColumnImpl.scala | 3 +- .../sql/catalyst/parser/DDLParserSuite.scala | 213 +++++++++++++++++- .../catalog/InMemoryTableCatalog.scala | 3 +- .../datasources/DataSourceStrategy.scala | 7 +- .../datasources/v2/DataSourceV2Strategy.scala | 5 +- .../sql-tests/results/ansi/keywords.sql.out | 2 + .../sql-tests/results/keywords.sql.out | 2 + .../sql/connector/DataSourceV2SQLSuite.scala | 58 +++++ .../sql/execution/command/DDLSuite.scala | 11 + .../ThriftServerWithSparkContextSuite.scala | 2 +- 22 files changed, 724 insertions(+), 41 deletions(-) create mode 100644 sql/api/src/main/java/org/apache/spark/sql/connector/catalog/IdentityColumnSpec.java create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/IdentityColumn.scala diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index a6d8550716b96..38472f44fb599 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1589,6 +1589,30 @@ ], "sqlState" : "42601" }, + "IDENTITY_COLUMNS_DUPLICATED_SEQUENCE_GENERATOR_OPTION" : { + "message" : [ + "Duplicated IDENTITY column sequence generator option: ." + ], + "sqlState" : "42601" + }, + "IDENTITY_COLUMNS_ILLEGAL_STEP" : { + "message" : [ + "IDENTITY column step cannot be 0." + ], + "sqlState" : "42611" + }, + "IDENTITY_COLUMNS_UNSUPPORTED_DATA_TYPE" : { + "message" : [ + "DataType is not supported for IDENTITY columns." + ], + "sqlState" : "428H2" + }, + "IDENTITY_COLUMN_WITH_DEFAULT_VALUE" : { + "message" : [ + "A column cannot have both a default value and an identity column specification but column has default value: () and identity column specification: ()." + ], + "sqlState" : "42623" + }, "ILLEGAL_DAY_OF_WEEK" : { "message" : [ "Illegal input for day of week: ." diff --git a/docs/sql-ref-ansi-compliance.md b/docs/sql-ref-ansi-compliance.md index fe5ddf27bf6c4..7987e5eb6012a 100644 --- a/docs/sql-ref-ansi-compliance.md +++ b/docs/sql-ref-ansi-compliance.md @@ -536,12 +536,14 @@ Below is a list of all the keywords in Spark SQL. |HOUR|non-reserved|non-reserved|non-reserved| |HOURS|non-reserved|non-reserved|non-reserved| |IDENTIFIER|non-reserved|non-reserved|non-reserved| +|IDENTITY|non-reserved|non-reserved|non-reserved| |IF|non-reserved|non-reserved|not a keyword| |IGNORE|non-reserved|non-reserved|non-reserved| |IMMEDIATE|non-reserved|non-reserved|non-reserved| |IMPORT|non-reserved|non-reserved|non-reserved| |IN|reserved|non-reserved|reserved| |INCLUDE|non-reserved|non-reserved|non-reserved| +|INCREMENT|non-reserved|non-reserved|non-reserved| |INDEX|non-reserved|non-reserved|non-reserved| |INDEXES|non-reserved|non-reserved|non-reserved| |INNER|reserved|strict-non-reserved|reserved| diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index 96a58b99debeb..c82ee57a25179 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -256,12 +256,14 @@ BINARY_HEX: 'X'; HOUR: 'HOUR'; HOURS: 'HOURS'; IDENTIFIER_KW: 'IDENTIFIER'; +IDENTITY: 'IDENTITY'; IF: 'IF'; IGNORE: 'IGNORE'; IMMEDIATE: 'IMMEDIATE'; IMPORT: 'IMPORT'; IN: 'IN'; INCLUDE: 'INCLUDE'; +INCREMENT: 'INCREMENT'; INDEX: 'INDEX'; INDEXES: 'INDEXES'; INNER: 'INNER'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 3ea408ca42703..1840b68878419 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -1297,7 +1297,22 @@ colDefinitionOption ; generationExpression - : GENERATED ALWAYS AS LEFT_PAREN expression RIGHT_PAREN + : GENERATED ALWAYS AS LEFT_PAREN expression RIGHT_PAREN #generatedColumn + | GENERATED (ALWAYS | BY DEFAULT) AS IDENTITY identityColSpec? #identityColumn + ; + +identityColSpec + : LEFT_PAREN sequenceGeneratorOption* RIGHT_PAREN + ; + +sequenceGeneratorOption + : START WITH start=sequenceGeneratorStartOrStep + | INCREMENT BY step=sequenceGeneratorStartOrStep + ; + +sequenceGeneratorStartOrStep + : MINUS? INTEGER_VALUE + | MINUS? BIGINT_LITERAL ; complexColTypeList @@ -1591,11 +1606,13 @@ ansiNonReserved | HOUR | HOURS | IDENTIFIER_KW + | IDENTITY | IF | IGNORE | IMMEDIATE | IMPORT | INCLUDE + | INCREMENT | INDEX | INDEXES | INPATH @@ -1942,12 +1959,14 @@ nonReserved | HOUR | HOURS | IDENTIFIER_KW + | IDENTITY | IF | IGNORE | IMMEDIATE | IMPORT | IN | INCLUDE + | INCREMENT | INDEX | INDEXES | INPATH diff --git a/sql/api/src/main/java/org/apache/spark/sql/connector/catalog/IdentityColumnSpec.java b/sql/api/src/main/java/org/apache/spark/sql/connector/catalog/IdentityColumnSpec.java new file mode 100644 index 0000000000000..4a8943736bd31 --- /dev/null +++ b/sql/api/src/main/java/org/apache/spark/sql/connector/catalog/IdentityColumnSpec.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.catalog; +import org.apache.spark.annotation.Evolving; + +import java.util.Objects; + +/** + * Identity column specification. + */ +@Evolving +public class IdentityColumnSpec { + private final long start; + private final long step; + private final boolean allowExplicitInsert; + + /** + * Creates an identity column specification. + * @param start the start value to generate the identity values + * @param step the step value to generate the identity values + * @param allowExplicitInsert whether the identity column allows explicit insertion of values + */ + public IdentityColumnSpec(long start, long step, boolean allowExplicitInsert) { + this.start = start; + this.step = step; + this.allowExplicitInsert = allowExplicitInsert; + } + + /** + * @return the start value to generate the identity values + */ + public long getStart() { + return start; + } + + /** + * @return the step value to generate the identity values + */ + public long getStep() { + return step; + } + + /** + * @return whether the identity column allows explicit insertion of values + */ + public boolean isAllowExplicitInsert() { + return allowExplicitInsert; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + IdentityColumnSpec that = (IdentityColumnSpec) o; + return start == that.start && + step == that.step && + allowExplicitInsert == that.allowExplicitInsert; + } + + @Override + public int hashCode() { + return Objects.hash(start, step, allowExplicitInsert); + } + + @Override + public String toString() { + return "IdentityColumnSpec{" + + "start=" + start + + ", step=" + step + + ", allowExplicitInsert=" + allowExplicitInsert + + "}"; + } +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala index 5f7fcb92f7bd1..b19607a28f06c 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala @@ -556,6 +556,25 @@ private[sql] object QueryParsingErrors extends DataTypeErrorsBase { ctx) } + def identityColumnUnsupportedDataType( + ctx: IdentityColumnContext, + dataType: String): Throwable = { + new ParseException("IDENTITY_COLUMNS_UNSUPPORTED_DATA_TYPE", Map("dataType" -> dataType), ctx) + } + + def identityColumnIllegalStep(ctx: IdentityColSpecContext): Throwable = { + new ParseException("IDENTITY_COLUMNS_ILLEGAL_STEP", Map.empty, ctx) + } + + def identityColumnDuplicatedSequenceGeneratorOption( + ctx: IdentityColSpecContext, + sequenceGeneratorOption: String): Throwable = { + new ParseException( + "IDENTITY_COLUMNS_DUPLICATED_SEQUENCE_GENERATOR_OPTION", + Map("sequenceGeneratorOption" -> sequenceGeneratorOption), + ctx) + } + def createViewWithBothIfNotExistsAndReplaceError(ctx: CreateViewContext): Throwable = { new ParseException(errorClass = "_LEGACY_ERROR_TEMP_0052", ctx) } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/Column.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/Column.java index b191438dbc3ee..8b32940d7a657 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/Column.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/Column.java @@ -53,7 +53,7 @@ static Column create( boolean nullable, String comment, String metadataInJSON) { - return new ColumnImpl(name, dataType, nullable, comment, null, null, metadataInJSON); + return new ColumnImpl(name, dataType, nullable, comment, null, null, null, metadataInJSON); } static Column create( @@ -63,7 +63,8 @@ static Column create( String comment, ColumnDefaultValue defaultValue, String metadataInJSON) { - return new ColumnImpl(name, dataType, nullable, comment, defaultValue, null, metadataInJSON); + return new ColumnImpl(name, dataType, nullable, comment, defaultValue, + null, null, metadataInJSON); } static Column create( @@ -74,7 +75,18 @@ static Column create( String generationExpression, String metadataInJSON) { return new ColumnImpl(name, dataType, nullable, comment, null, - generationExpression, metadataInJSON); + generationExpression, null, metadataInJSON); + } + + static Column create( + String name, + DataType dataType, + boolean nullable, + String comment, + IdentityColumnSpec identityColumnSpec, + String metadataInJSON) { + return new ColumnImpl(name, dataType, nullable, comment, null, + null, identityColumnSpec, metadataInJSON); } /** @@ -113,6 +125,12 @@ static Column create( @Nullable String generationExpression(); + /** + * Returns the identity column specification of this table column. Null means no identity column. + */ + @Nullable + IdentityColumnSpec identityColumnSpec(); + /** * Returns the column metadata in JSON format. */ diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java index 5ccb15ff1f0a4..dceac1b484cf2 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java @@ -59,5 +59,23 @@ public enum TableCatalogCapability { * {@link TableCatalog#createTable}. * See {@link Column#defaultValue()}. */ - SUPPORT_COLUMN_DEFAULT_VALUE + SUPPORT_COLUMN_DEFAULT_VALUE, + + /** + * Signals that the TableCatalog supports defining identity columns upon table creation in SQL. + *

+ * Without this capability, any create/replace table statements with an identity column defined + * in the table schema will throw an exception during analysis. + *

+ * An identity column is defined with syntax: + * {@code colName colType GENERATED ALWAYS AS IDENTITY(identityColumnSpec)} + * or + * {@code colName colType GENERATED BY DEFAULT AS IDENTITY(identityColumnSpec)} + * identityColumnSpec is defined with syntax: {@code [START WITH start | INCREMENT BY step]*} + *

+ * IdentitySpec is included in the column definition for APIs like + * {@link TableCatalog#createTable}. + * See {@link Column#identityColumnSpec()}. + */ + SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index edcb417da123b..cb0e0e35c3704 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -45,7 +45,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.PARAMETER import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, DateTimeUtils, IntervalUtils} import org.apache.spark.sql.catalyst.util.DateTimeUtils.{convertSpecialDate, convertSpecialTimestamp, convertSpecialTimestampNTZ, getZoneId, stringToDate, stringToTimestamp, stringToTimestampWithoutTimeZone} -import org.apache.spark.sql.connector.catalog.{CatalogV2Util, SupportsNamespaces, TableCatalog, TableWritePrivilege} +import org.apache.spark.sql.connector.catalog.{CatalogV2Util, IdentityColumnSpec, SupportsNamespaces, TableCatalog, TableWritePrivilege} import org.apache.spark.sql.connector.catalog.TableChange.ColumnPosition import org.apache.spark.sql.connector.expressions.{ApplyTransform, BucketTransform, DaysTransform, Expression => V2Expression, FieldReference, HoursTransform, IdentityTransform, LiteralValue, MonthsTransform, Transform, YearsTransform} import org.apache.spark.sql.errors.{DataTypeErrorsBase, QueryCompilationErrors, QueryParsingErrors, SqlScriptingErrors} @@ -3619,13 +3619,19 @@ class AstBuilder extends DataTypeAstBuilder } } + val dataType = typedVisit[DataType](ctx.dataType) ColumnDefinition( name = name, - dataType = typedVisit[DataType](ctx.dataType), + dataType = dataType, nullable = nullable, comment = commentSpec.map(visitCommentSpec), defaultValue = defaultExpression.map(visitDefaultExpression), - generationExpression = generationExpression.map(visitGenerationExpression) + generationExpression = generationExpression.collect { + case ctx: GeneratedColumnContext => visitGeneratedColumn(ctx) + }, + identityColumnSpec = generationExpression.collect { + case ctx: IdentityColumnContext => visitIdentityColumn(ctx, dataType) + } ) } @@ -3681,11 +3687,63 @@ class AstBuilder extends DataTypeAstBuilder /** * Create a generation expression string. */ - override def visitGenerationExpression(ctx: GenerationExpressionContext): String = + override def visitGeneratedColumn(ctx: GeneratedColumnContext): String = withOrigin(ctx) { getDefaultExpression(ctx.expression(), "GENERATED").originalSQL } + /** + * Parse and verify IDENTITY column definition. + * + * @param ctx The parser context. + * @param dataType The data type of column defined as IDENTITY column. Used for verification. + * @return Tuple containing start, step and allowExplicitInsert. + */ + protected def visitIdentityColumn( + ctx: IdentityColumnContext, + dataType: DataType): IdentityColumnSpec = { + if (dataType != LongType && dataType != IntegerType) { + throw QueryParsingErrors.identityColumnUnsupportedDataType(ctx, dataType.toString) + } + // We support two flavors of syntax: + // (1) GENERATED ALWAYS AS IDENTITY (...) + // (2) GENERATED BY DEFAULT AS IDENTITY (...) + // (1) forbids explicit inserts, while (2) allows. + val allowExplicitInsert = ctx.BY() != null && ctx.DEFAULT() != null + val (start, step) = visitIdentityColSpec(ctx.identityColSpec()) + + new IdentityColumnSpec(start, step, allowExplicitInsert) + } + + override def visitIdentityColSpec(ctx: IdentityColSpecContext): (Long, Long) = { + val defaultStart = 1 + val defaultStep = 1 + if (ctx == null) { + return (defaultStart, defaultStep) + } + var (start, step): (Option[Long], Option[Long]) = (None, None) + ctx.sequenceGeneratorOption().asScala.foreach { option => + if (option.start != null) { + if (start.isDefined) { + throw QueryParsingErrors.identityColumnDuplicatedSequenceGeneratorOption(ctx, "START") + } + start = Some(option.start.getText.toLong) + } else if (option.step != null) { + if (step.isDefined) { + throw QueryParsingErrors.identityColumnDuplicatedSequenceGeneratorOption(ctx, "STEP") + } + step = Some(option.step.getText.toLong) + if (step.get == 0L) { + throw QueryParsingErrors.identityColumnIllegalStep(ctx) + } + } else { + throw SparkException + .internalError(s"Invalid identity column sequence generator option: ${option.getText}") + } + } + (start.getOrElse(defaultStart), step.getOrElse(defaultStep)) + } + /** * Create an optional comment string. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ColumnDefinition.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ColumnDefinition.scala index 83e50aa33c70d..043214711ccf9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ColumnDefinition.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ColumnDefinition.scala @@ -21,10 +21,10 @@ import org.apache.spark.SparkException import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.expressions.{Expression, Literal, UnaryExpression, Unevaluable} import org.apache.spark.sql.catalyst.parser.ParserInterface -import org.apache.spark.sql.catalyst.util.GeneratedColumn +import org.apache.spark.sql.catalyst.util.{GeneratedColumn, IdentityColumn} import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.validateDefaultValueExpr import org.apache.spark.sql.catalyst.util.ResolveDefaultColumnsUtils.{CURRENT_DEFAULT_COLUMN_METADATA_KEY, EXISTS_DEFAULT_COLUMN_METADATA_KEY} -import org.apache.spark.sql.connector.catalog.{Column => V2Column, ColumnDefaultValue} +import org.apache.spark.sql.connector.catalog.{Column => V2Column, ColumnDefaultValue, IdentityColumnSpec} import org.apache.spark.sql.connector.expressions.LiteralValue import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.connector.ColumnImpl @@ -41,7 +41,11 @@ case class ColumnDefinition( comment: Option[String] = None, defaultValue: Option[DefaultValueExpression] = None, generationExpression: Option[String] = None, + identityColumnSpec: Option[IdentityColumnSpec] = None, metadata: Metadata = Metadata.empty) extends Expression with Unevaluable { + assert( + generationExpression.isEmpty || identityColumnSpec.isEmpty, + "A ColumnDefinition cannot contain both a generation expression and an identity column spec.") override def children: Seq[Expression] = defaultValue.toSeq @@ -58,6 +62,7 @@ case class ColumnDefinition( comment.orNull, defaultValue.map(_.toV2(statement, name)).orNull, generationExpression.orNull, + identityColumnSpec.orNull, if (metadata == Metadata.empty) null else metadata.json) } @@ -75,8 +80,19 @@ case class ColumnDefinition( generationExpression.foreach { generationExpr => metadataBuilder.putString(GeneratedColumn.GENERATION_EXPRESSION_METADATA_KEY, generationExpr) } + encodeIdentityColumnSpec(metadataBuilder) StructField(name, dataType, nullable, metadataBuilder.build()) } + + private def encodeIdentityColumnSpec(metadataBuilder: MetadataBuilder): Unit = { + identityColumnSpec.foreach { spec: IdentityColumnSpec => + metadataBuilder.putLong(IdentityColumn.IDENTITY_INFO_START, spec.getStart) + metadataBuilder.putLong(IdentityColumn.IDENTITY_INFO_STEP, spec.getStep) + metadataBuilder.putBoolean( + IdentityColumn.IDENTITY_INFO_ALLOW_EXPLICIT_INSERT, + spec.isAllowExplicitInsert) + } + } } object ColumnDefinition { @@ -87,6 +103,9 @@ object ColumnDefinition { metadataBuilder.remove(CURRENT_DEFAULT_COLUMN_METADATA_KEY) metadataBuilder.remove(EXISTS_DEFAULT_COLUMN_METADATA_KEY) metadataBuilder.remove(GeneratedColumn.GENERATION_EXPRESSION_METADATA_KEY) + metadataBuilder.remove(IdentityColumn.IDENTITY_INFO_START) + metadataBuilder.remove(IdentityColumn.IDENTITY_INFO_STEP) + metadataBuilder.remove(IdentityColumn.IDENTITY_INFO_ALLOW_EXPLICIT_INSERT) val hasDefaultValue = col.getCurrentDefaultValue().isDefined && col.getExistenceDefaultValue().isDefined @@ -97,6 +116,15 @@ object ColumnDefinition { None } val generationExpr = GeneratedColumn.getGenerationExpression(col) + val identityColumnSpec = if (col.metadata.contains(IdentityColumn.IDENTITY_INFO_START)) { + Some(new IdentityColumnSpec( + col.metadata.getLong(IdentityColumn.IDENTITY_INFO_START), + col.metadata.getLong(IdentityColumn.IDENTITY_INFO_STEP), + col.metadata.getBoolean(IdentityColumn.IDENTITY_INFO_ALLOW_EXPLICIT_INSERT) + )) + } else { + None + } ColumnDefinition( col.name, col.dataType, @@ -104,6 +132,7 @@ object ColumnDefinition { col.getComment(), defaultValue, generationExpr, + identityColumnSpec, metadataBuilder.build() ) } @@ -124,18 +153,8 @@ object ColumnDefinition { s"Command $cmd should not have column default value expression.") } cmd.columns.foreach { col => - if (col.defaultValue.isDefined && col.generationExpression.isDefined) { - throw new AnalysisException( - errorClass = "GENERATED_COLUMN_WITH_DEFAULT_VALUE", - messageParameters = Map( - "colName" -> col.name, - "defaultValue" -> col.defaultValue.get.originalSQL, - "genExpr" -> col.generationExpression.get - ) - ) - } - col.defaultValue.foreach { default => + checkDefaultColumnConflicts(col) validateDefaultValueExpr(default, statement, col.name, col.dataType) } } @@ -143,6 +162,29 @@ object ColumnDefinition { case _ => } } + + private def checkDefaultColumnConflicts(col: ColumnDefinition): Unit = { + if (col.generationExpression.isDefined) { + throw new AnalysisException( + errorClass = "GENERATED_COLUMN_WITH_DEFAULT_VALUE", + messageParameters = Map( + "colName" -> col.name, + "defaultValue" -> col.defaultValue.get.originalSQL, + "genExpr" -> col.generationExpression.get + ) + ) + } + if (col.identityColumnSpec.isDefined) { + throw new AnalysisException( + errorClass = "IDENTITY_COLUMN_WITH_DEFAULT_VALUE", + messageParameters = Map( + "colName" -> col.name, + "defaultValue" -> col.defaultValue.get.originalSQL, + "identityColumnSpec" -> col.identityColumnSpec.get.toString + ) + ) + } + } } /** diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/IdentityColumn.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/IdentityColumn.scala new file mode 100644 index 0000000000000..26a3cb026d317 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/IdentityColumn.scala @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.util + +import org.apache.spark.sql.connector.catalog.{Identifier, IdentityColumnSpec, TableCatalog, TableCatalogCapability} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.types.{StructField, StructType} + +/** + * This object contains utility methods and values for Identity Columns + */ +object IdentityColumn { + val IDENTITY_INFO_START = "identity.start" + val IDENTITY_INFO_STEP = "identity.step" + val IDENTITY_INFO_ALLOW_EXPLICIT_INSERT = "identity.allowExplicitInsert" + + /** + * If `schema` contains any generated columns, check whether the table catalog supports identity + * columns. Otherwise throw an error. + */ + def validateIdentityColumn( + schema: StructType, + catalog: TableCatalog, + ident: Identifier): Unit = { + if (hasIdentityColumns(schema)) { + if (!catalog + .capabilities() + .contains(TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + catalog, ident, operation = "identity column" + ) + } + } + } + + /** + * Whether the given `field` is an identity column + */ + def isIdentityColumn(field: StructField): Boolean = { + field.metadata.contains(IDENTITY_INFO_START) + } + + /** + * Returns the identity information stored in the column metadata if it exists + */ + def getIdentityInfo(field: StructField): Option[IdentityColumnSpec] = { + if (isIdentityColumn(field)) { + Some(new IdentityColumnSpec( + field.metadata.getString(IDENTITY_INFO_START).toLong, + field.metadata.getString(IDENTITY_INFO_STEP).toLong, + field.metadata.getString(IDENTITY_INFO_ALLOW_EXPLICIT_INSERT).toBoolean)) + } else { + None + } + } + + /** + * Whether the `schema` has one or more identity columns + */ + def hasIdentityColumns(schema: StructType): Boolean = { + schema.exists(isIdentityColumn) + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala index 6698f0a021400..9b7f68070a1a4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.analysis.{AsOfTimestamp, AsOfVersion, Named import org.apache.spark.sql.catalyst.catalog.ClusterBySpec import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.catalyst.plans.logical.{SerdeInfo, TableSpec} -import org.apache.spark.sql.catalyst.util.GeneratedColumn +import org.apache.spark.sql.catalyst.util.{GeneratedColumn, IdentityColumn} import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns._ import org.apache.spark.sql.connector.catalog.TableChange._ import org.apache.spark.sql.connector.catalog.functions.UnboundFunction @@ -579,18 +579,10 @@ private[sql] object CatalogV2Util { val isDefaultColumn = f.getCurrentDefaultValue().isDefined && f.getExistenceDefaultValue().isDefined val isGeneratedColumn = GeneratedColumn.isGeneratedColumn(f) - if (isDefaultColumn && isGeneratedColumn) { - throw new AnalysisException( - errorClass = "GENERATED_COLUMN_WITH_DEFAULT_VALUE", - messageParameters = Map( - "colName" -> f.name, - "defaultValue" -> f.getCurrentDefaultValue().get, - "genExpr" -> GeneratedColumn.getGenerationExpression(f).get - ) - ) - } - + val isIdentityColumn = IdentityColumn.isIdentityColumn(f) if (isDefaultColumn) { + checkDefaultColumnConflicts(f) + val e = analyze( f, statementType = "Column analysis", @@ -611,10 +603,41 @@ private[sql] object CatalogV2Util { Seq("comment", GeneratedColumn.GENERATION_EXPRESSION_METADATA_KEY)) Column.create(f.name, f.dataType, f.nullable, f.getComment().orNull, GeneratedColumn.getGenerationExpression(f).get, metadataAsJson(cleanedMetadata)) + } else if (isIdentityColumn) { + val cleanedMetadata = metadataWithKeysRemoved( + Seq("comment", + IdentityColumn.IDENTITY_INFO_START, + IdentityColumn.IDENTITY_INFO_STEP, + IdentityColumn.IDENTITY_INFO_ALLOW_EXPLICIT_INSERT)) + Column.create(f.name, f.dataType, f.nullable, f.getComment().orNull, + IdentityColumn.getIdentityInfo(f).get, metadataAsJson(cleanedMetadata)) } else { val cleanedMetadata = metadataWithKeysRemoved(Seq("comment")) Column.create(f.name, f.dataType, f.nullable, f.getComment().orNull, metadataAsJson(cleanedMetadata)) } } + + private def checkDefaultColumnConflicts(f: StructField): Unit = { + if (GeneratedColumn.isGeneratedColumn(f)) { + throw new AnalysisException( + errorClass = "GENERATED_COLUMN_WITH_DEFAULT_VALUE", + messageParameters = Map( + "colName" -> f.name, + "defaultValue" -> f.getCurrentDefaultValue().get, + "genExpr" -> GeneratedColumn.getGenerationExpression(f).get + ) + ) + } + if (IdentityColumn.isIdentityColumn(f)) { + throw new AnalysisException( + errorClass = "IDENTITY_COLUMN_WITH_DEFAULT_VALUE", + messageParameters = Map( + "colName" -> f.name, + "defaultValue" -> f.getCurrentDefaultValue().get, + "identityColumnSpec" -> IdentityColumn.getIdentityInfo(f).get.toString + ) + ) + } + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ColumnImpl.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ColumnImpl.scala index 2a67ffc4bbef5..47889410561e3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ColumnImpl.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ColumnImpl.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.internal.connector -import org.apache.spark.sql.connector.catalog.{Column, ColumnDefaultValue} +import org.apache.spark.sql.connector.catalog.{Column, ColumnDefaultValue, IdentityColumnSpec} import org.apache.spark.sql.types.DataType // The standard concrete implementation of data source V2 column. @@ -28,4 +28,5 @@ case class ColumnImpl( comment: String, defaultValue: ColumnDefaultValue, generationExpression: String, + identityColumnSpec: IdentityColumnSpec, metadataInJSON: String) extends Column diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala index 0f2bb791f3465..b7e2490b552cc 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala @@ -20,14 +20,16 @@ package org.apache.spark.sql.catalyst.parser import java.util.Locale import org.apache.spark.SparkThrowable +import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.expressions.{EqualTo, Hex, Literal} import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.connector.catalog.IdentityColumnSpec import org.apache.spark.sql.connector.catalog.TableChange.ColumnPosition.{after, first} import org.apache.spark.sql.connector.expressions.{ApplyTransform, BucketTransform, ClusterByTransform, DaysTransform, FieldReference, HoursTransform, IdentityTransform, LiteralValue, MonthsTransform, Transform, YearsTransform} import org.apache.spark.sql.connector.expressions.LogicalExpressions.bucket import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{Decimal, IntegerType, LongType, StringType, StructType, TimestampType} +import org.apache.spark.sql.types.{DataType, Decimal, IntegerType, LongType, StringType, StructType, TimestampType} import org.apache.spark.storage.StorageLevelMapper import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} @@ -2856,10 +2858,217 @@ class DDLParserSuite extends AnalysisTest { exception = parseException( "CREATE TABLE my_tab(a INT, b INT GENERATED ALWAYS AS a + 1) USING PARQUET"), condition = "PARSE_SYNTAX_ERROR", - parameters = Map("error" -> "'a'", "hint" -> ": missing '('") + parameters = Map("error" -> "'a'", "hint" -> "") ) } + test("SPARK-48824: implement parser support for " + + "GENERATED ALWAYS/BY DEFAULT AS IDENTITY columns in tables ") { + def parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr: String, + identityColumnDefStr: String, + identityColumnSpecStr: String, + expectedDataType: DataType, + expectedStart: Long, + expectedStep: Long, + expectedAllowExplicitInsert: Boolean): Unit = { + val columnsWithIdentitySpec = Seq( + ColumnDefinition( + name = "id", + dataType = expectedDataType, + nullable = true, + identityColumnSpec = Some( + new IdentityColumnSpec( + expectedStart, + expectedStep, + expectedAllowExplicitInsert + ) + ) + ), + ColumnDefinition("val", IntegerType) + ) + comparePlans( + parsePlan( + s"CREATE TABLE my_tab(id $identityColumnDataTypeStr GENERATED $identityColumnDefStr" + + s" AS IDENTITY $identityColumnSpecStr, val INT) USING parquet" + ), + CreateTable( + UnresolvedIdentifier(Seq("my_tab")), + columnsWithIdentitySpec, + Seq.empty[Transform], + UnresolvedTableSpec( + Map.empty[String, String], + Some("parquet"), + OptionList(Seq.empty), + None, + None, + None, + false + ), + false + ) + ) + + comparePlans( + parsePlan( + s"REPLACE TABLE my_tab(id $identityColumnDataTypeStr GENERATED $identityColumnDefStr" + + s" AS IDENTITY $identityColumnSpecStr, val INT) USING parquet" + ), + ReplaceTable( + UnresolvedIdentifier(Seq("my_tab")), + columnsWithIdentitySpec, + Seq.empty[Transform], + UnresolvedTableSpec( + Map.empty[String, String], + Some("parquet"), + OptionList(Seq.empty), + None, + None, + None, + false + ), + false + ) + ) + } + for { + identityColumnDefStr <- Seq("BY DEFAULT", "ALWAYS") + identityColumnDataTypeStr <- Seq("BIGINT", "INT") + } { + val expectedAllowExplicitInsert = identityColumnDefStr == "BY DEFAULT" + val expectedDataType = identityColumnDataTypeStr match { + case "BIGINT" => LongType + case "INT" => IntegerType + } + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(START WITH 2 INCREMENT BY 2)", + expectedDataType, + expectedStart = 2, + expectedStep = 2, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(START WITH -2 INCREMENT BY -2)", + expectedDataType, + expectedStart = -2, + expectedStep = -2, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(START WITH 2)", + expectedDataType, + expectedStart = 2, + expectedStep = 1, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(START WITH -2)", + expectedDataType, + expectedStart = -2, + expectedStep = 1, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(INCREMENT BY 2)", + expectedDataType, + expectedStart = 1, + expectedStep = 2, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "(INCREMENT BY -2)", + expectedDataType, + expectedStart = 1, + expectedStep = -2, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "()", + expectedDataType, + expectedStart = 1, + expectedStep = 1, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + parseAndCompareIdentityColumnPlan( + identityColumnDataTypeStr, + identityColumnDefStr, + "", + expectedDataType, + expectedStart = 1, + expectedStep = 1, + expectedAllowExplicitInsert = expectedAllowExplicitInsert) + } + } + + test("SPARK-48824: Column cannot have both a generation expression and an identity column spec") { + checkError( + exception = intercept[AnalysisException] { + parsePlan(s"CREATE TABLE testcat.my_tab(id BIGINT GENERATED ALWAYS AS 1" + + s" GENERATED ALWAYS AS IDENTITY, val INT) USING foo") + }, + condition = "PARSE_SYNTAX_ERROR", + parameters = Map("error" -> "'1'", "hint" -> "") + ) + } + + test("SPARK-48824: Identity column step must not be zero") { + checkError( + exception = intercept[ParseException] { + parsePlan( + s"CREATE TABLE testcat.my_tab" + + s"(id BIGINT GENERATED ALWAYS AS IDENTITY(INCREMENT BY 0), val INT) USING foo" + ) + }, + condition = "IDENTITY_COLUMNS_ILLEGAL_STEP", + parameters = Map.empty, + context = ExpectedContext( + fragment = "id BIGINT GENERATED ALWAYS AS IDENTITY(INCREMENT BY 0)", + start = 28, + stop = 81) + ) + } + + test("SPARK-48824: Identity column datatype must be long or integer") { + checkError( + exception = intercept[ParseException] { + parsePlan( + s"CREATE TABLE testcat.my_tab(id FLOAT GENERATED ALWAYS AS IDENTITY(), val INT) USING foo" + ) + }, + condition = "IDENTITY_COLUMNS_UNSUPPORTED_DATA_TYPE", + parameters = Map("dataType" -> "FloatType"), + context = + ExpectedContext(fragment = "id FLOAT GENERATED ALWAYS AS IDENTITY()", start = 28, stop = 66) + ) + } + + test("SPARK-48824: Identity column sequence generator option cannot be duplicated") { + val identityColumnSpecStrs = Seq( + "(START WITH 0 START WITH 1)", + "(INCREMENT BY 1 INCREMENT BY 2)", + "(START WITH 0 INCREMENT BY 1 START WITH 1)", + "(INCREMENT BY 1 START WITH 0 INCREMENT BY 2)" + ) + for { + identitySpecStr <- identityColumnSpecStrs + } { + val exception = intercept[ParseException] { + parsePlan( + s"CREATE TABLE testcat.my_tab" + + s"(id BIGINT GENERATED ALWAYS AS IDENTITY $identitySpecStr, val INT) USING foo" + ) + } + assert(exception.getErrorClass === "IDENTITY_COLUMNS_DUPLICATED_SEQUENCE_GENERATOR_OPTION") + } + } + test("SPARK-42681: Relax ordering constraint for ALTER TABLE ADD COLUMN options") { // Positive test cases to verify that column definition options could be applied in any order. val expectedPlan = AddColumns( diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala index 982de88e58847..56ed3bb243e19 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala @@ -167,7 +167,8 @@ class InMemoryTableCatalog extends BasicInMemoryTableCatalog with SupportsNamesp override def capabilities: java.util.Set[TableCatalogCapability] = { Set( TableCatalogCapability.SUPPORT_COLUMN_DEFAULT_VALUE, - TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_GENERATED_COLUMNS + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_GENERATED_COLUMNS, + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS ).asJava } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala index 1dd2659a1b169..2be4b236872f0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{AppendData, InsertIntoDir, I import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.catalyst.util.{GeneratedColumn, ResolveDefaultColumns, V2ExpressionBuilder} +import org.apache.spark.sql.catalyst.util.{GeneratedColumn, IdentityColumn, ResolveDefaultColumns, V2ExpressionBuilder} import org.apache.spark.sql.connector.catalog.{SupportsRead, V1Table} import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, NullOrdering, SortDirection, SortOrder => V2SortOrder, SortValue} @@ -146,6 +146,11 @@ object DataSourceAnalysis extends Rule[LogicalPlan] { tableDesc.identifier, "generated columns") } + if (IdentityColumn.hasIdentityColumns(newSchema)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + tableDesc.identifier, "identity columns") + } + val newTableDesc = tableDesc.copy(schema = newSchema) CreateDataSourceTableCommand(newTableDesc, ignoreIfExists = mode == SaveMode.Ignore) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index 112ee2c5450b2..d7f46c32f99a0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -32,7 +32,8 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, DynamicPruning import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.util.{toPrettySQL, GeneratedColumn, ResolveDefaultColumns, V2ExpressionBuilder} +import org.apache.spark.sql.catalyst.util.{toPrettySQL, GeneratedColumn, + IdentityColumn, ResolveDefaultColumns, V2ExpressionBuilder} import org.apache.spark.sql.connector.catalog.{Identifier, StagingTableCatalog, SupportsDeleteV2, SupportsNamespaces, SupportsPartitionManagement, SupportsWrite, Table, TableCapability, TableCatalog, TruncatableTable} import org.apache.spark.sql.connector.catalog.index.SupportsIndex import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue} @@ -185,6 +186,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat val statementType = "CREATE TABLE" GeneratedColumn.validateGeneratedColumns( c.tableSchema, catalog.asTableCatalog, ident, statementType) + IdentityColumn.validateIdentityColumn(c.tableSchema, catalog.asTableCatalog, ident) CreateTableExec( catalog.asTableCatalog, @@ -214,6 +216,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat val statementType = "REPLACE TABLE" GeneratedColumn.validateGeneratedColumns( c.tableSchema, catalog.asTableCatalog, ident, statementType) + IdentityColumn.validateIdentityColumn(c.tableSchema, catalog.asTableCatalog, ident) val v2Columns = columns.map(_.toV2Column(statementType)).toArray catalog match { diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out index 81ccc0f9efc13..b464427d379a3 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out @@ -142,6 +142,7 @@ HAVING true HOUR false HOURS false IDENTIFIER false +IDENTITY false IF false IGNORE false ILIKE false @@ -149,6 +150,7 @@ IMMEDIATE false IMPORT false IN true INCLUDE false +INCREMENT false INDEX false INDEXES false INNER true diff --git a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out index e145c57332eb2..16436d7a722ce 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out @@ -142,6 +142,7 @@ HAVING false HOUR false HOURS false IDENTIFIER false +IDENTITY false IF false IGNORE false ILIKE false @@ -149,6 +150,7 @@ IMMEDIATE false IMPORT false IN false INCLUDE false +INCREMENT false INDEX false INDEXES false INNER false diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index 998d459cd436c..5df7b62cfb285 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -1753,6 +1753,64 @@ class DataSourceV2SQLSuiteV1Filter } } + test("SPARK-48824: Column cannot have both an identity column spec and a default value") { + val tblName = "my_tab" + val tableDefinition = + s"$tblName(id BIGINT GENERATED ALWAYS AS IDENTITY DEFAULT 0, name STRING)" + withSQLConf(SQLConf.DEFAULT_COLUMN_ALLOWED_PROVIDERS.key -> "foo") { + for (statement <- Seq("CREATE TABLE", "REPLACE TABLE")) { + withTable(s"testcat.$tblName") { + if (statement == "REPLACE TABLE") { + sql(s"CREATE TABLE testcat.$tblName(a INT) USING foo") + } + checkError( + exception = intercept[AnalysisException] { + sql(s"$statement testcat.$tableDefinition USING foo") + }, + condition = "IDENTITY_COLUMN_WITH_DEFAULT_VALUE", + parameters = Map( + "colName" -> "id", + "defaultValue" -> "0", + "identityColumnSpec" -> + "IdentityColumnSpec{start=1, step=1, allowExplicitInsert=false}") + ) + } + } + } + } + + test("SPARK-48824: Identity columns only allowed with TableCatalogs that " + + "SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS") { + val tblName = "my_tab" + val tableDefinition = + s"$tblName(id BIGINT GENERATED ALWAYS AS IDENTITY(), val INT)" + for (statement <- Seq("CREATE TABLE", "REPLACE TABLE")) { + // InMemoryTableCatalog.capabilities() = {SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS} + withTable(s"testcat.$tblName") { + if (statement == "REPLACE TABLE") { + sql(s"CREATE TABLE testcat.$tblName(a INT) USING foo") + } + // Can create table with an identity column + sql(s"$statement testcat.$tableDefinition USING foo") + assert(catalog("testcat").asTableCatalog.tableExists(Identifier.of(Array(), tblName))) + } + // BasicInMemoryTableCatalog.capabilities() = {} + withSQLConf("spark.sql.catalog.dummy" -> classOf[BasicInMemoryTableCatalog].getName) { + checkError( + exception = intercept[AnalysisException] { + sql("USE dummy") + sql(s"$statement dummy.$tableDefinition USING foo") + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`dummy`.`my_tab`", + "operation" -> "identity column" + ) + ) + } + } + } + test("SPARK-46972: asymmetrical replacement for char/varchar in V2SessionCatalog.createTable") { // unset this config to use the default v2 session catalog. spark.conf.unset(V2_SESSION_CATALOG_IMPLEMENTATION.key) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala index 6e58b0e62ed63..8307326f17fcf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala @@ -2288,6 +2288,17 @@ abstract class DDLSuite extends QueryTest with DDLSuiteBase { ) } + test("SPARK-48824: No identity columns with V1") { + checkError( + exception = intercept[AnalysisException] { + sql(s"create table t(a int, b bigint generated always as identity()) using parquet") + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map("tableName" -> "`spark_catalog`.`default`.`t`", + "operation" -> "identity columns") + ) + } + test("SPARK-44837: Error when altering partition column in non-delta table") { withTable("t") { sql("CREATE TABLE t(i INT, j INT, k INT) USING parquet PARTITIONED BY (i, j)") diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala index edef6371be8ae..5b8ee4ea9714f 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala @@ -214,7 +214,7 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer { val sessionHandle = client.openSession(user, "") val infoValue = client.getInfo(sessionHandle, GetInfoType.CLI_ODBC_KEYWORDS) // scalastyle:off line.size.limit - assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") + assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") // scalastyle:on line.size.limit } } From b4f4d9b7a7d470158af39d75824bcc501e3506da Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Sun, 15 Sep 2024 18:17:34 -0700 Subject: [PATCH 036/189] [SPARK-49655][BUILD] Link `python3` to `python3.9` in `spark-rm` Docker image ### What changes were proposed in this pull request? This PR aims to link `python3` to `python3.9` in `spark-rm` docker image. ### Why are the changes needed? We already link `python` to `python3.9`. https://github.com/apache/spark/blob/931ab065df3952487028316ebd49c2895d947bf2/dev/create-release/spark-rm/Dockerfile#L139 We need to link `python3` to `python3.9` to fix Spark Documentation generation failure in release script. ``` $ dev/create-release/do-release-docker.sh -d /run/user/1000/spark -s docs ... = Building documentation... Command: /opt/spark-rm/release-build.sh docs Log file: docs.log Command FAILED. Check full logs for details. from /opt/spark-rm/output/spark/docs/.local_ruby_bundle/ruby/3.0.0/gems/jekyll-4.3.3/lib/jekyll/command.rb:91:in `process_with_graceful_fail' ``` The root cause is `mkdocs` module import error during `error-conditions.html` generation. ### Does this PR introduce _any_ user-facing change? No. This is a release-script. ### How was this patch tested? Manual review. After this PR, `error docs` generation succeeds. ``` ************************ * Building error docs. * ************************ Generated: docs/_generated/error-conditions.html ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48117 from dongjoon-hyun/SPARK-49655. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- dev/create-release/spark-rm/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/dev/create-release/spark-rm/Dockerfile b/dev/create-release/spark-rm/Dockerfile index e7f558b523d0c..3cba72d042ed6 100644 --- a/dev/create-release/spark-rm/Dockerfile +++ b/dev/create-release/spark-rm/Dockerfile @@ -137,6 +137,7 @@ RUN python3.9 -m pip list RUN gem install --no-document "bundler:2.4.22" RUN ln -s "$(which python3.9)" "/usr/local/bin/python" +RUN ln -s "$(which python3.9)" "/usr/local/bin/python3" WORKDIR /opt/spark-rm/output From 738db079c0b65e8305b7a1349923ee017316f691 Mon Sep 17 00:00:00 2001 From: Avery Qi Date: Mon, 16 Sep 2024 11:37:37 +0800 Subject: [PATCH 037/189] [SPARK-49646][SQL] fix subquery decorrelation for union/set operations when parentOuterReferences has references not covered in collectedChildOuterReferences ### What changes were proposed in this pull request? fix bug when encounter union/setOp under limit/aggregation with filter predicates cannot pulled up directly in lateral join. eg: ``` create table IF NOT EXISTS t(t1 INT,t2 int) using json; CREATE TABLE IF NOT EXISTS a (a1 INT) using json; select 1 from t as t_outer left join lateral( select b1,b2 from ( select a.a1 as b1, 1 as b2 from a union select t_outer.t1 as b1, null as b2 ) as t_inner where (t_inner.b1 < t_outer.t2 or t_inner.b1 is null) and t_inner.b1 = t_outer.t1 order by t_inner.b1,t_inner.b2 desc limit 1 ) as lateral_table ``` ### Why are the changes needed? In general, spark cannot handle this query because: 1. Decorrelation logic tries to rewrite limit operator into Window aggregation and pull up correlated predicates, and Union operator is rewritten to have DomainJoin within its children with outer references. 2. When we're rewriting DomainJoin to real join execution, it needs attribute reference map based on pulled up correlated predicates to rewrite outer references in DomainJoin. However, each child of Union/SetOp operator are using different attribute references even they are referring to the same column of outer table. We need Union/SetOp output and its children output to map between these references. 3. Combined with aggregation and filters with inequality comparison, more outer references are remained within children of Union operator, and these references are not covered in Union/SetOp output which leads to lacking of information when we're trying to map different attributed references within children of Union/SetOp operator. More context -> please read this short investigation doc(I've changed the link and it's now public): https://docs.google.com/document/d/1_pJIi_8GuLHOXabLEgRy2e7OHw-OIBnWbwGwSkwIcxg/edit?usp=sharing ### Does this PR introduce _any_ user-facing change? yes, bug is fixed and the above query can be handled without error. ### How was this patch tested? added unit test ### Was this patch authored or co-authored using generative AI tooling? No Closes #48109 from averyqi-db/averyqi-db/SPARK-49646. Authored-by: Avery Qi Signed-off-by: Wenchen Fan --- .../optimizer/DecorrelateInnerQuery.scala | 2 +- .../analyzer-results/join-lateral.sql.out | 47 +++++++++++++++++++ .../sql-tests/inputs/join-lateral.sql | 21 +++++++++ .../sql-tests/results/join-lateral.sql.out | 27 +++++++++++ 4 files changed, 96 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala index 424f4b96271d3..6c0d7189862d3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala @@ -1064,7 +1064,7 @@ object DecorrelateInnerQuery extends PredicateHelper { // Project, they could get added at the beginning or the end of the output columns // depending on the child plan. // The inner expressions for the domain are the values of newOuterReferenceMap. - val domainProjections = collectedChildOuterReferences.map(newOuterReferenceMap(_)) + val domainProjections = newOuterReferences.map(newOuterReferenceMap(_)) val newChild = Project(child.output ++ domainProjections, decorrelatedChild) (newChild, newJoinCond, newOuterReferenceMap) } diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/join-lateral.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/join-lateral.sql.out index e81ee769f57d6..5bf893605423c 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/join-lateral.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/join-lateral.sql.out @@ -3017,6 +3017,53 @@ Project [c1#x, c2#x, t#x] +- LocalRelation [col1#x, col2#x] +-- !query +select 1 +from t1 as t_outer +left join + lateral( + select b1,b2 + from + ( + select + t2.c1 as b1, + 1 as b2 + from t2 + union + select t_outer.c1 as b1, + null as b2 + ) as t_inner + where (t_inner.b1 < t_outer.c2 or t_inner.b1 is null) + and t_inner.b1 = t_outer.c1 + order by t_inner.b1,t_inner.b2 desc limit 1 + ) as lateral_table +-- !query analysis +Project [1 AS 1#x] ++- LateralJoin lateral-subquery#x [c2#x && c1#x && c1#x], LeftOuter + : +- SubqueryAlias lateral_table + : +- GlobalLimit 1 + : +- LocalLimit 1 + : +- Sort [b1#x ASC NULLS FIRST, b2#x DESC NULLS LAST], true + : +- Project [b1#x, b2#x] + : +- Filter (((b1#x < outer(c2#x)) OR isnull(b1#x)) AND (b1#x = outer(c1#x))) + : +- SubqueryAlias t_inner + : +- Distinct + : +- Union false, false + : :- Project [c1#x AS b1#x, 1 AS b2#x] + : : +- SubqueryAlias spark_catalog.default.t2 + : : +- View (`spark_catalog`.`default`.`t2`, [c1#x, c2#x]) + : : +- Project [cast(col1#x as int) AS c1#x, cast(col2#x as int) AS c2#x] + : : +- LocalRelation [col1#x, col2#x] + : +- Project [b1#x, cast(b2#x as int) AS b2#x] + : +- Project [outer(c1#x) AS b1#x, null AS b2#x] + : +- OneRowRelation + +- SubqueryAlias t_outer + +- SubqueryAlias spark_catalog.default.t1 + +- View (`spark_catalog`.`default`.`t1`, [c1#x, c2#x]) + +- Project [cast(col1#x as int) AS c1#x, cast(col2#x as int) AS c2#x] + +- LocalRelation [col1#x, col2#x] + + -- !query DROP VIEW t1 -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/join-lateral.sql b/sql/core/src/test/resources/sql-tests/inputs/join-lateral.sql index 8bff1f109aa65..e3cef9207d20f 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/join-lateral.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/join-lateral.sql @@ -531,6 +531,27 @@ select * from t1 join lateral (select t4.c1 as t from t4 where t1.c1 = t4.c1)) as foo order by foo.t limit 5); + +select 1 +from t1 as t_outer +left join + lateral( + select b1,b2 + from + ( + select + t2.c1 as b1, + 1 as b2 + from t2 + union + select t_outer.c1 as b1, + null as b2 + ) as t_inner + where (t_inner.b1 < t_outer.c2 or t_inner.b1 is null) + and t_inner.b1 = t_outer.c1 + order by t_inner.b1,t_inner.b2 desc limit 1 + ) as lateral_table; + -- clean up DROP VIEW t1; DROP VIEW t2; diff --git a/sql/core/src/test/resources/sql-tests/results/join-lateral.sql.out b/sql/core/src/test/resources/sql-tests/results/join-lateral.sql.out index ced8d6398a66f..11bafb2cf63c9 100644 --- a/sql/core/src/test/resources/sql-tests/results/join-lateral.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/join-lateral.sql.out @@ -1878,6 +1878,33 @@ struct 1 2 3 +-- !query +select 1 +from t1 as t_outer +left join + lateral( + select b1,b2 + from + ( + select + t2.c1 as b1, + 1 as b2 + from t2 + union + select t_outer.c1 as b1, + null as b2 + ) as t_inner + where (t_inner.b1 < t_outer.c2 or t_inner.b1 is null) + and t_inner.b1 = t_outer.c1 + order by t_inner.b1,t_inner.b2 desc limit 1 + ) as lateral_table +-- !query schema +struct<1:int> +-- !query output +1 +1 + + -- !query DROP VIEW t1 -- !query schema From 2113f109b8d73cb8deb404664f25bd51308ca809 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Mon, 16 Sep 2024 16:33:44 +0800 Subject: [PATCH 038/189] [SPARK-49611][SQL] Introduce TVF `collations()` & remove the `SHOW COLLATIONS` command ### What changes were proposed in this pull request? The pr aims to - introduce `TVF` `collations()`. - remove the `SHOW COLLATIONS` command. ### Why are the changes needed? Based on cloud-fan's suggestion: https://github.com/apache/spark/pull/47364#issuecomment-2345183501 I believe that after this, we can do many things based on it, such as `filtering` and `querying` based on `LANGUAGE` or `COUNTRY`, etc. eg: ```sql SELECT * FROM collations() WHERE LANGUAGE like '%Chinese%'; ``` ### Does this PR introduce _any_ user-facing change? Yes, provide a new TVF `collations()` for end-users. ### How was this patch tested? - Add new UT. - Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48087 from panbingkun/SPARK-49611. Lead-authored-by: panbingkun Co-authored-by: panbingkun Signed-off-by: Wenchen Fan --- docs/sql-ref-ansi-compliance.md | 1 - .../spark/sql/catalyst/parser/SqlBaseLexer.g4 | 1 - .../sql/catalyst/parser/SqlBaseParser.g4 | 2 - .../catalyst/analysis/FunctionRegistry.scala | 1 + .../sql/catalyst/catalog/SessionCatalog.scala | 15 +--- .../sql/catalyst/expressions/generators.scala | 44 ++++++++++- .../ansi-sql-2016-reserved-keywords.txt | 1 - .../spark/sql/execution/SparkSqlParser.scala | 12 --- .../command/ShowCollationsCommand.scala | 62 --------------- .../sql-tests/results/ansi/keywords.sql.out | 2 - .../sql-tests/results/keywords.sql.out | 1 - .../org/apache/spark/sql/CollationSuite.scala | 79 +++++++++++++------ .../ThriftServerWithSparkContextSuite.scala | 2 +- 13 files changed, 101 insertions(+), 122 deletions(-) delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala diff --git a/docs/sql-ref-ansi-compliance.md b/docs/sql-ref-ansi-compliance.md index 7987e5eb6012a..fff6906457f7d 100644 --- a/docs/sql-ref-ansi-compliance.md +++ b/docs/sql-ref-ansi-compliance.md @@ -442,7 +442,6 @@ Below is a list of all the keywords in Spark SQL. |CODEGEN|non-reserved|non-reserved|non-reserved| |COLLATE|reserved|non-reserved|reserved| |COLLATION|reserved|non-reserved|reserved| -|COLLATIONS|reserved|non-reserved|reserved| |COLLECTION|non-reserved|non-reserved|non-reserved| |COLUMN|reserved|non-reserved|reserved| |COLUMNS|non-reserved|non-reserved|non-reserved| diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index c82ee57a25179..e704f9f58b964 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -162,7 +162,6 @@ CLUSTERED: 'CLUSTERED'; CODEGEN: 'CODEGEN'; COLLATE: 'COLLATE'; COLLATION: 'COLLATION'; -COLLATIONS: 'COLLATIONS'; COLLECTION: 'COLLECTION'; COLUMN: 'COLUMN'; COLUMNS: 'COLUMNS'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 1840b68878419..f13dde773496a 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -268,7 +268,6 @@ statement | SHOW PARTITIONS identifierReference partitionSpec? #showPartitions | SHOW identifier? FUNCTIONS ((FROM | IN) ns=identifierReference)? (LIKE? (legacy=multipartIdentifier | pattern=stringLit))? #showFunctions - | SHOW COLLATIONS (LIKE? pattern=stringLit)? #showCollations | SHOW CREATE TABLE identifierReference (AS SERDE)? #showCreateTable | SHOW CURRENT namespace #showCurrentNamespace | SHOW CATALOGS (LIKE? pattern=stringLit)? #showCatalogs @@ -1868,7 +1867,6 @@ nonReserved | CODEGEN | COLLATE | COLLATION - | COLLATIONS | COLLECTION | COLUMN | COLUMNS diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala index 75e1ab86f1772..5a3c4b0ec8696 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala @@ -1158,6 +1158,7 @@ object TableFunctionRegistry { generator[PosExplode]("posexplode"), generator[PosExplode]("posexplode_outer", outer = true), generator[Stack]("stack"), + generator[Collations]("collations"), generator[SQLKeywords]("sql_keywords"), generator[VariantExplode]("variant_explode"), generator[VariantExplode]("variant_explode_outer", outer = true) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala index 5c14e261fafc8..d3a6cb6ae2845 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala @@ -24,7 +24,6 @@ import java.util.concurrent.TimeUnit import javax.annotation.concurrent.GuardedBy import scala.collection.mutable -import scala.jdk.CollectionConverters.CollectionHasAsScala import scala.util.{Failure, Success, Try} import com.google.common.cache.{Cache, CacheBuilder} @@ -40,8 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Expression, Expre import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParserInterface} import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, SubqueryAlias, View} import org.apache.spark.sql.catalyst.trees.CurrentOrigin -import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, CollationFactory, StringUtils} -import org.apache.spark.sql.catalyst.util.CollationFactory.CollationMeta +import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, StringUtils} import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} @@ -1901,17 +1899,6 @@ class SessionCatalog( .filter(isTemporaryFunction) } - /** - * List all built-in collations with the given pattern. - */ - def listCollations(pattern: Option[String]): Seq[CollationMeta] = { - val collationIdentifiers = CollationFactory.listCollations().asScala.toSeq - val filteredCollationNames = StringUtils.filterPattern( - collationIdentifiers.map(_.getName), pattern.getOrElse("*")).toSet - collationIdentifiers.filter(ident => filteredCollationNames.contains(ident.getName)).map( - CollationFactory.loadCollationMeta) - } - // ----------------- // | Other methods | // ----------------- diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala index 2cc88a25f465d..dc58352a1b362 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.catalyst.expressions import scala.collection.mutable +import scala.jdk.CollectionConverters.CollectionHasAsScala import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} @@ -28,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.plans.logical.{FunctionSignature, InputParameter} import org.apache.spark.sql.catalyst.trees.TreePattern.{GENERATOR, TreePattern} -import org.apache.spark.sql.catalyst.util.{ArrayData, MapData} +import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory, MapData} import org.apache.spark.sql.catalyst.util.SQLKeywordUtils._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf @@ -618,3 +619,44 @@ case class SQLKeywords() extends LeafExpression with Generator with CodegenFallb override def prettyName: String = "sql_keywords" } + +@ExpressionDescription( + usage = """_FUNC_() - Get all of the Spark SQL string collations""", + examples = """ + Examples: + > SELECT * FROM _FUNC_() WHERE NAME = 'UTF8_BINARY'; + SYSTEM BUILTIN UTF8_BINARY NULL NULL ACCENT_SENSITIVE CASE_SENSITIVE NO_PAD NULL + """, + since = "4.0.0", + group = "generator_funcs") +case class Collations() extends LeafExpression with Generator with CodegenFallback { + override def elementSchema: StructType = new StructType() + .add("CATALOG", StringType, nullable = false) + .add("SCHEMA", StringType, nullable = false) + .add("NAME", StringType, nullable = false) + .add("LANGUAGE", StringType) + .add("COUNTRY", StringType) + .add("ACCENT_SENSITIVITY", StringType, nullable = false) + .add("CASE_SENSITIVITY", StringType, nullable = false) + .add("PAD_ATTRIBUTE", StringType, nullable = false) + .add("ICU_VERSION", StringType) + + override def eval(input: InternalRow): IterableOnce[InternalRow] = { + CollationFactory.listCollations().asScala.map(CollationFactory.loadCollationMeta).map { m => + InternalRow( + UTF8String.fromString(m.catalog), + UTF8String.fromString(m.schema), + UTF8String.fromString(m.collationName), + UTF8String.fromString(m.language), + UTF8String.fromString(m.country), + UTF8String.fromString( + if (m.accentSensitivity) "ACCENT_SENSITIVE" else "ACCENT_INSENSITIVE"), + UTF8String.fromString( + if (m.caseSensitivity) "CASE_SENSITIVE" else "CASE_INSENSITIVE"), + UTF8String.fromString(m.padAttribute), + UTF8String.fromString(m.icuVersion)) + } + } + + override def prettyName: String = "collations" +} diff --git a/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt b/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt index 452cf930525bc..46da60b7897b8 100644 --- a/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt +++ b/sql/catalyst/src/test/resources/ansi-sql-2016-reserved-keywords.txt @@ -48,7 +48,6 @@ CLOSE COALESCE COLLATE COLLATION -COLLATIONS COLLECT COLUMN COMMIT diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala index 640abaea58abe..a8261e5d98ba0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala @@ -1096,16 +1096,4 @@ class SparkSqlAstBuilder extends AstBuilder { withIdentClause(ctx.identifierReference(), UnresolvedNamespace(_)), cleanedProperties) } - - /** - * Create a [[ShowCollationsCommand]] command. - * Expected format: - * {{{ - * SHOW COLLATIONS (LIKE? pattern=stringLit)?; - * }}} - */ - override def visitShowCollations(ctx: ShowCollationsContext): LogicalPlan = withOrigin(ctx) { - val pattern = Option(ctx.pattern).map(x => string(visitStringLit(x))) - ShowCollationsCommand(pattern) - } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala deleted file mode 100644 index 179a841b013bd..0000000000000 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ShowCollationsCommand.scala +++ /dev/null @@ -1,62 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.spark.sql.execution.command - -import org.apache.spark.sql.{Row, SparkSession} -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} -import org.apache.spark.sql.catalyst.util.CollationFactory.CollationMeta -import org.apache.spark.sql.types.StringType - -/** - * A command for `SHOW COLLATIONS`. - * - * The syntax of this command is: - * {{{ - * SHOW COLLATIONS (LIKE? pattern=stringLit)?; - * }}} - */ -case class ShowCollationsCommand(pattern: Option[String]) extends LeafRunnableCommand { - - override val output: Seq[Attribute] = Seq( - AttributeReference("COLLATION_CATALOG", StringType, nullable = false)(), - AttributeReference("COLLATION_SCHEMA", StringType, nullable = false)(), - AttributeReference("COLLATION_NAME", StringType, nullable = false)(), - AttributeReference("LANGUAGE", StringType)(), - AttributeReference("COUNTRY", StringType)(), - AttributeReference("ACCENT_SENSITIVITY", StringType, nullable = false)(), - AttributeReference("CASE_SENSITIVITY", StringType, nullable = false)(), - AttributeReference("PAD_ATTRIBUTE", StringType, nullable = false)(), - AttributeReference("ICU_VERSION", StringType)()) - - override def run(sparkSession: SparkSession): Seq[Row] = { - val systemCollations: Seq[CollationMeta] = - sparkSession.sessionState.catalog.listCollations(pattern) - - systemCollations.map(m => Row( - m.catalog, - m.schema, - m.collationName, - m.language, - m.country, - if (m.accentSensitivity) "ACCENT_SENSITIVE" else "ACCENT_INSENSITIVE", - if (m.caseSensitivity) "CASE_SENSITIVE" else "CASE_INSENSITIVE", - m.padAttribute, - m.icuVersion - )) - } -} diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out index b464427d379a3..6497a46c68ccd 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out @@ -48,7 +48,6 @@ CLUSTERED false CODEGEN false COLLATE true COLLATION true -COLLATIONS true COLLECTION false COLUMN true COLUMNS false @@ -384,7 +383,6 @@ CAST CHECK COLLATE COLLATION -COLLATIONS COLUMN CONSTRAINT CREATE diff --git a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out index 16436d7a722ce..0dfd62599afa6 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out @@ -48,7 +48,6 @@ CLUSTERED false CODEGEN false COLLATE false COLLATION false -COLLATIONS false COLLECTION false COLUMN false COLUMNS false diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index b25cddb80762a..489a990d3e1cf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -1625,38 +1625,38 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { } } - test("show collations") { - assert(sql("SHOW COLLATIONS").collect().length >= 562) + test("TVF collations()") { + assert(sql("SELECT * FROM collations()").collect().length >= 562) // verify that the output ordering is as expected (UTF8_BINARY, UTF8_LCASE, etc.) - val df = sql("SHOW COLLATIONS").limit(10) + val df = sql("SELECT * FROM collations() limit 10") checkAnswer(df, Seq(Row("SYSTEM", "BUILTIN", "UTF8_BINARY", null, null, "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", null), - Row("SYSTEM", "BUILTIN", "UTF8_LCASE", null, null, - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", null), - Row("SYSTEM", "BUILTIN", "UNICODE", "", "", - "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "UNICODE_AI", "", "", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "UNICODE_CI", "", "", - "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "UNICODE_CI_AI", "", "", - "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "af", "Afrikaans", "", - "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "af_AI", "Afrikaans", "", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "af_CI", "Afrikaans", "", - "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "af_CI_AI", "Afrikaans", "", - "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) - - checkAnswer(sql("SHOW COLLATIONS LIKE '*UTF8_BINARY*'"), + Row("SYSTEM", "BUILTIN", "UTF8_LCASE", null, null, + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", null), + Row("SYSTEM", "BUILTIN", "UNICODE", "", "", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_AI", "", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_CI", "", "", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_CI_AI", "", "", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af", "Afrikaans", "", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_AI", "Afrikaans", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_CI", "Afrikaans", "", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_CI_AI", "Afrikaans", "", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + + checkAnswer(sql("SELECT * FROM collations() WHERE NAME LIKE '%UTF8_BINARY%'"), Row("SYSTEM", "BUILTIN", "UTF8_BINARY", null, null, "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", null)) - checkAnswer(sql("SHOW COLLATIONS '*zh_Hant_HKG*'"), + checkAnswer(sql("SELECT * FROM collations() WHERE NAME LIKE '%zh_Hant_HKG%'"), Seq(Row("SYSTEM", "BUILTIN", "zh_Hant_HKG", "Chinese", "Hong Kong SAR China", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_AI", "Chinese", "Hong Kong SAR China", @@ -1665,5 +1665,36 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI_AI", "Chinese", "Hong Kong SAR China", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + + checkAnswer(sql("SELECT * FROM collations() WHERE COUNTRY = 'Singapore'"), + Seq(Row("SYSTEM", "BUILTIN", "zh_Hans_SGP", "Chinese", "Singapore", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_AI", "Chinese", "Singapore", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_CI", "Chinese", "Singapore", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_CI_AI", "Chinese", "Singapore", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + + checkAnswer(sql("SELECT * FROM collations() WHERE LANGUAGE = 'English' " + + "and COUNTRY = 'United States'"), + Seq(Row("SYSTEM", "BUILTIN", "en_USA", "English", "United States", + "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "en_USA_AI", "English", "United States", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "en_USA_CI", "English", "United States", + "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "en_USA_CI_AI", "English", "United States", + "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) + + checkAnswer(sql("SELECT NAME, LANGUAGE, ACCENT_SENSITIVITY, CASE_SENSITIVITY " + + "FROM collations() WHERE COUNTRY = 'United States'"), + Seq(Row("en_USA", "English", "ACCENT_SENSITIVE", "CASE_SENSITIVE"), + Row("en_USA_AI", "English", "ACCENT_SENSITIVE", "CASE_INSENSITIVE"), + Row("en_USA_CI", "English", "ACCENT_INSENSITIVE", "CASE_SENSITIVE"), + Row("en_USA_CI_AI", "English", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE"))) + + checkAnswer(sql("SELECT NAME FROM collations() WHERE ICU_VERSION is null"), + Seq(Row("UTF8_BINARY"), Row("UTF8_LCASE"))) } } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala index 5b8ee4ea9714f..4bc4116a23da7 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala @@ -214,7 +214,7 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer { val sessionHandle = client.openSession(user, "") val infoValue = client.getInfo(sessionHandle, GetInfoType.CLI_ODBC_KEYWORDS) // scalastyle:off line.size.limit - assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") + assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") // scalastyle:on line.size.limit } } From 8d78f5b02c256423ee125c31bf746a1a15dbcf25 Mon Sep 17 00:00:00 2001 From: Enrico Minack Date: Mon, 16 Sep 2024 12:30:04 -0400 Subject: [PATCH 039/189] [DOCS][MINOR] Remove spark-sql-api dependency from connect docs ### What changes were proposed in this pull request? Remove `spark-sql-api` dependency from documentation. ### Why are the changes needed? Dependency `spark-connect-client-jvm` is sufficient, as it includes (shaded) the `spark-sql-api` package. In fact, adding the `spark-sql-api` dependency breaks runtime: 1. transient dependency `io.netty:netty-buffer`, that is also included in `spark-connect-client-jvm` (shaded) cannot be found: `NoClassDefFoundError: io/netty/buffer/PooledByteBufAllocator` 2. method `ArrowUtils$.toArrowSchema` provided by `spark-sql-api` cannot be found: `NoSuchMethodError: 'org.sparkproject.org.apache.arrow.vector.types.pojo.Schema org.apache.spark.sql.util.ArrowUtils$.toArrowSchema(org.apache.spark.sql.types.StructType, java.lang.String, boolean, boolean)'` ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Manually with minimal maven project. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48119 from EnricoMi/docs-connect-jvm-deps. Authored-by: Enrico Minack Signed-off-by: Herman van Hovell --- docs/spark-connect-overview.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index b77f71fb695db..1cc409bfbc007 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -335,7 +335,6 @@ Lines with a: 72, lines with b: 39 To use Spark Connect as part of a Scala application/project, we first need to include the right dependencies. Using the `sbt` build system as an example, we add the following dependencies to the `build.sbt` file: {% highlight sbt %} -libraryDependencies += "org.apache.spark" %% "spark-sql-api" % "3.5.0" libraryDependencies += "org.apache.spark" %% "spark-connect-client-jvm" % "3.5.0" {% endhighlight %} From 294af6e31639d6f6ac51f961f319866f077b5302 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Mon, 16 Sep 2024 20:52:28 -0700 Subject: [PATCH 040/189] [SPARK-49680][PYTHON] Limit `Sphinx` build parallelism to 4 by default ### What changes were proposed in this pull request? This PR aims to limit `Sphinx` build parallelism to 4 by default for the following goals. - This will preserve the same speed in GitHub Action environment. - This will prevent the exhaustive `SparkSubmit` invocation in large machines like `c6i.24xlarge`. - The user still can override by providing `SPHINXOPTS`. ### Why are the changes needed? `Sphinx` parallelism feature was added via the following on 2024-01-10. - #44680 However, unfortunately, this breaks Python API doc generation in large machines because this means the number of parallel `SparkSubmit` invocation of PySpark. In addition, given that each `PySpark` currently is launched with `local[*]`, this ends up `N * N` `pyspark.daemon`s. In other words, as of today, this default setting, `auto`, seems to work on low-core machine like `GitHub Action` runners (4 cores). For example, this breaks `Python` documentations build even on M3 Max environment and this is worse on large EC2 machines (c7i.24xlarge). You can see the failure locally like this. ``` $ build/sbt package -Phive-thriftserver $ cd python/docs $ make html ... 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4041. Attempting port 4042. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4042. Attempting port 4043. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4041. Attempting port 4042. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4042. Attempting port 4043. 24/09/16 17:04:38 WARN Utils: Service 'SparkUI' could not bind on port 4043. Attempting port 4044. 24/09/16 17:04:39 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041. 24/09/16 17:04:39 WARN Utils: Service 'SparkUI' could not bind on port 4041. Attempting port 4042. 24/09/16 17:04:39 WARN Utils: Service 'SparkUI' could not bind on port 4042. Attempting port 4043. 24/09/16 17:04:39 WARN Utils: Service 'SparkUI' could not bind on port 4043. Attempting port 4044. 24/09/16 17:04:39 WARN Utils: Service 'SparkUI' could not bind on port 4044. Attempting port 4045. ... java.lang.OutOfMemoryError: Java heap space ... 24/09/16 14:09:55 WARN PythonRunner: Incomplete task 7.0 in stage 30 (TID 177) interrupted: Attempting to kill Python Worker ... make: *** [html] Error 2 ``` ### Does this PR introduce _any_ user-facing change? No, this is a dev-only change. ### How was this patch tested? Pass the CIs and do manual tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48129 from dongjoon-hyun/SPARK-49680. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- python/docs/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/docs/Makefile b/python/docs/Makefile index 5058c1206171b..428b0d24b568e 100644 --- a/python/docs/Makefile +++ b/python/docs/Makefile @@ -16,7 +16,7 @@ # Minimal makefile for Sphinx documentation # You can set these variables from the command line. -SPHINXOPTS ?= "-W" "-j" "auto" +SPHINXOPTS ?= "-W" "-j" "4" SPHINXBUILD ?= sphinx-build SOURCEDIR ?= source BUILDDIR ?= build From 370453adba1730b5412750b34e87a35147d71aa2 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Mon, 16 Sep 2024 20:53:35 -0700 Subject: [PATCH 041/189] [SPARK-49678][CORE] Support `spark.test.master` in `SparkSubmitArguments` ### What changes were proposed in this pull request? This PR aims to support `spark.test.master` in `SparkSubmitArguments`. ### Why are the changes needed? To allow users to control the default master setting during testing and documentation generation. #### First, currently, we cannot build `Python Documentation` on M3 Max (and high-core machines) without this. Only it succeeds on GitHub Action runners (4 cores) or equivalent low-core docker run. Please try the following on your Macs. **BEFORE** ``` $ build/sbt package -Phive-thriftserver $ cd python/docs $ make html ... java.lang.OutOfMemoryError: Java heap space ... 24/09/16 14:09:55 WARN PythonRunner: Incomplete task 7.0 in stage 30 (TID 177) interrupted: Attempting to kill Python Worker ... make: *** [html] Error 2 ``` **AFTER** ``` $ build/sbt package -Phive-thriftserver $ cd python/docs $ JDK_JAVA_OPTIONS="-Dspark.test.master=local[1]" make html ... build succeeded. The HTML pages are in build/html. ``` #### Second, in general, we can control all `SparkSubmit` (eg. Spark Shells) like the following. **BEFORE (`local[*]`)** ``` $ bin/pyspark Python 3.9.19 (main, Jun 17 2024, 15:39:29) [Clang 15.0.0 (clang-1500.3.9.4)] on darwin Type "help", "copyright", "credits" or "license" for more information. WARNING: Using incubator modules: jdk.incubator.vector Using Spark's default log4j profile: org/apache/spark/log4j2-pattern-layout-defaults.properties Setting default log level to "WARN". To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel). 24/09/16 13:53:02 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable Welcome to ____ __ / __/__ ___ _____/ /__ _\ \/ _ \/ _ `/ __/ '_/ /__ / .__/\_,_/_/ /_/\_\ version 4.0.0-SNAPSHOT /_/ Using Python version 3.9.19 (main, Jun 17 2024 15:39:29) Spark context Web UI available at http://localhost:4040 Spark context available as 'sc' (master = local[*], app id = local-1726519982935). SparkSession available as 'spark'. >>> ``` **AFTER (`local[1]`)** ``` $ JDK_JAVA_OPTIONS="-Dspark.test.master=local[1]" bin/pyspark NOTE: Picked up JDK_JAVA_OPTIONS: -Dspark.test.master=local[1] Python 3.9.19 (main, Jun 17 2024, 15:39:29) [Clang 15.0.0 (clang-1500.3.9.4)] on darwin Type "help", "copyright", "credits" or "license" for more information. NOTE: Picked up JDK_JAVA_OPTIONS: -Dspark.test.master=local[1] NOTE: Picked up JDK_JAVA_OPTIONS: -Dspark.test.master=local[1] WARNING: Using incubator modules: jdk.incubator.vector Using Spark's default log4j profile: org/apache/spark/log4j2-pattern-layout-defaults.properties Setting default log level to "WARN". To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel). 24/09/16 13:51:03 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable Welcome to ____ __ / __/__ ___ _____/ /__ _\ \/ _ \/ _ `/ __/ '_/ /__ / .__/\_,_/_/ /_/\_\ version 4.0.0-SNAPSHOT /_/ Using Python version 3.9.19 (main, Jun 17 2024 15:39:29) Spark context Web UI available at http://localhost:4040 Spark context available as 'sc' (master = local[1], app id = local-1726519863363). SparkSession available as 'spark'. >>> ``` ### Does this PR introduce _any_ user-facing change? No. `spark.test.master` is a new parameter. ### How was this patch tested? Manual tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48126 from dongjoon-hyun/SPARK-49678. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../scala/org/apache/spark/deploy/SparkSubmitArguments.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmitArguments.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmitArguments.scala index 32dd2f81bbc82..2c9ddff348056 100644 --- a/core/src/main/scala/org/apache/spark/deploy/SparkSubmitArguments.scala +++ b/core/src/main/scala/org/apache/spark/deploy/SparkSubmitArguments.scala @@ -43,7 +43,8 @@ private[deploy] class SparkSubmitArguments(args: Seq[String], env: Map[String, S extends SparkSubmitArgumentsParser with Logging { var maybeMaster: Option[String] = None // Global defaults. These should be keep to minimum to avoid confusing behavior. - def master: String = maybeMaster.getOrElse("local[*]") + def master: String = + maybeMaster.getOrElse(System.getProperty("spark.test.master", "local[*]")) var maybeRemote: Option[String] = None var deployMode: String = null var executorMemory: String = null From db8468105bdbd0b6dd8722e169fe39b13a3ee44f Mon Sep 17 00:00:00 2001 From: Vladimir Golubev Date: Tue, 17 Sep 2024 08:15:23 +0200 Subject: [PATCH 042/189] [SPARK-49659][SQL] Add a nice user-facing error for scalar subqueries inside VALUES clause ### What changes were proposed in this pull request? Introduce a new `SCALAR_SUBQUERY_IN_VALUES` error, since we don't support scalar subqueries in the VALUES clause for now. ### Why are the changes needed? To make Spark user experience nicer. ### Does this PR introduce _any_ user-facing change? Yes, the exception type/message will be more descriptive ### How was this patch tested? - New unit test - Existing unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #48118 from vladimirg-db/vladimirg-db/introduce-user-facing-exception-for-inline-tables-with-scalar-subqueries. Authored-by: Vladimir Golubev Signed-off-by: Max Gekk --- .../main/resources/error/error-conditions.json | 5 +++++ .../sql/catalyst/analysis/CheckAnalysis.scala | 12 ++++++++++++ .../spark/sql/errors/QueryCompilationErrors.scala | 8 ++++++++ .../org/apache/spark/sql/SQLQuerySuite.scala | 15 +++++++++++++++ 4 files changed, 40 insertions(+) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 38472f44fb599..57b3d33741e98 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -5330,6 +5330,11 @@ "" ] }, + "SCALAR_SUBQUERY_IN_VALUES" : { + "message" : [ + "Scalar subqueries in the VALUES clause." + ] + }, "UNSUPPORTED_CORRELATED_EXPRESSION_IN_JOIN_CONDITION" : { "message" : [ "Correlated subqueries in the join predicate cannot reference both join inputs:", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala index a9fbe548ba39e..752ff49e1f90d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala @@ -250,6 +250,9 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB context = u.origin.getQueryContext, summary = u.origin.context.summary) + case u: UnresolvedInlineTable if unresolvedInlineTableContainsScalarSubquery(u) => + throw QueryCompilationErrors.inlineTableContainsScalarSubquery(u) + case command: V2PartitionCommand => command.table match { case r @ ResolvedTable(_, _, table, _) => table match { @@ -1559,6 +1562,15 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB case _ => } } + + private def unresolvedInlineTableContainsScalarSubquery( + unresolvedInlineTable: UnresolvedInlineTable) = { + unresolvedInlineTable.rows.exists { row => + row.exists { expression => + expression.exists(_.isInstanceOf[ScalarSubquery]) + } + } + } } // a heap of the preempted error that only keeps the top priority element, representing the sole diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index f1f8be3d15751..f268ef85ef1dd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -4112,4 +4112,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat "expr" -> expr.toString), origin = expr.origin) } + + def inlineTableContainsScalarSubquery(inlineTable: LogicalPlan): Throwable = { + new AnalysisException( + errorClass = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.SCALAR_SUBQUERY_IN_VALUES", + messageParameters = Map.empty, + origin = inlineTable.origin + ) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index b7d0039446f30..9beceda263797 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -4909,6 +4909,21 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark ) } } + + test("SPARK-49659: Unsupported scalar subqueries in VALUES") { + checkError( + exception = intercept[AnalysisException]( + sql("SELECT * FROM VALUES ((SELECT 1) + (SELECT 2))") + ), + condition = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.SCALAR_SUBQUERY_IN_VALUES", + parameters = Map(), + context = ExpectedContext( + fragment = "VALUES ((SELECT 1) + (SELECT 2))", + start = 14, + stop = 45 + ) + ) + } } case class Foo(bar: Option[String]) From dd8d1270679ce6cc99c217eca7e5ecba463fd7ab Mon Sep 17 00:00:00 2001 From: Josh Rosen Date: Tue, 17 Sep 2024 09:17:45 +0200 Subject: [PATCH 043/189] [SPARK-49650][SQL][DOCS] Updates references to deprecated Hive JDBC client params in tests and docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? This PR updates unit tests and docs to update references to long-deprecated Hive JDBC client connection string parameters. The current docs and some test code are using parameters that have been deprecated for nearly a decade since https://issues.apache.org/jira/browse/HIVE-6972 / https://github.com/apache/hive/commit/07082e8a851cb44a95dcae50bc84fb43cb1e84c6 , so I think it's safe to clean up the usages now. ### Why are the changes needed? While looking at some hive-thriftserver unit tests logs, I saw repeated spam of ``` 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: ***** JDBC param deprecation ***** 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: The use of hive.server2.transport.mode is deprecated. 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: Please use transportMode like so: jdbc:hive2://:/dbName;transportMode= 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: ***** JDBC param deprecation ***** 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: The use of hive.server2.thrift.http.path is deprecated. 20/06/05 06:35:55.442 pool-1-thread-1 WARN Utils: Please use httpPath like so: jdbc:hive2://:/dbName;httpPath= ``` ### Does this PR introduce _any_ user-facing change? No, it's just a test + documentation change, recommending syntax which has been long-supported (some tests were [already using the new parameters](https://github.com/apache/spark/blob/d3eb99f79e508d62fdb7e9bc595f0240ac021df5/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SharedThriftServer.scala#L81-L85)). ### How was this patch tested? Existing tests (let's wait to confirm that they pass in CI). ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48114 from JoshRosen/avoid-using-deprecated-hive-params. Lead-authored-by: Josh Rosen Co-authored-by: Josh Rosen Signed-off-by: Max Gekk --- docs/sql-distributed-sql-engine.md | 4 ++-- .../sql/hive/thriftserver/HiveThriftServer2Suites.scala | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/sql-distributed-sql-engine.md b/docs/sql-distributed-sql-engine.md index 734723f8c6235..ae8fd9c7211bd 100644 --- a/docs/sql-distributed-sql-engine.md +++ b/docs/sql-distributed-sql-engine.md @@ -83,7 +83,7 @@ Use the following setting to enable HTTP mode as system property or in `hive-sit To test, use beeline to connect to the JDBC/ODBC server in http mode with: - beeline> !connect jdbc:hive2://:/?hive.server2.transport.mode=http;hive.server2.thrift.http.path= + beeline> !connect jdbc:hive2://:/;transportMode=http;httpPath= If you closed a session and do CTAS, you must set `fs.%s.impl.disable.cache` to true in `hive-site.xml`. See more details in [[SPARK-21067]](https://issues.apache.org/jira/browse/SPARK-21067). @@ -94,4 +94,4 @@ To use the Spark SQL command line interface (CLI) from the shell: ./bin/spark-sql -For details, please refer to [Spark SQL CLI](sql-distributed-sql-engine-spark-sql-cli.html) \ No newline at end of file +For details, please refer to [Spark SQL CLI](sql-distributed-sql-engine-spark-sql-cli.html) diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala index 3ccbd23b71c98..4575549005f33 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala @@ -1430,9 +1430,9 @@ abstract class HiveThriftServer2TestBase extends SparkFunSuite with BeforeAndAft protected def jdbcUri(database: String = "default"): String = if (mode == ServerMode.http) { s"""jdbc:hive2://$localhost:$serverPort/ - |$database? - |hive.server2.transport.mode=http; - |hive.server2.thrift.http.path=cliservice; + |$database; + |transportMode=http; + |httpPath=cliservice;? |${hiveConfList}#${hiveVarList} """.stripMargin.split("\n").mkString.trim } else { From f586ffbf47eec2b578ced1139ff13415a709a4d5 Mon Sep 17 00:00:00 2001 From: Marko Date: Tue, 17 Sep 2024 11:01:39 +0200 Subject: [PATCH 044/189] [SPARK-49498][SQL][TESTS] Check differences between SR_AI and SR_Latn_AI collations ### What changes were proposed in this pull request? Added test to check difference between SR_AI and SR_Latn_AI collations. ### Why are the changes needed? Better testing. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Test added to `CollationSuite.scala` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47958 from ilicmarkodb/add_test_for_difference_between_SR_AI_and_SR_Latn_AI_collations. Authored-by: Marko Signed-off-by: Max Gekk --- .../org/apache/spark/sql/CollationSuite.scala | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index 489a990d3e1cf..d5d18b1ab081c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -186,6 +186,29 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { } } + test("check difference betweeen SR_AI and SR_Latn_AI collations") { + // scalastyle:off nonascii + Seq( + ("c", "ć"), + ("c", "č"), + ("ć", "č"), + ("C", "Ć"), + ("C", "Č"), + ("Ć", "Č"), + ("s", "š"), + ("S", "Š"), + ("z", "ž"), + ("Z", "Ž") + ).foreach { + case (c1, c2) => + // SR_Latn_AI + checkAnswer(sql(s"SELECT '$c1' = '$c2' COLLATE SR_Latn_AI"), Row(false)) + // SR_AI + checkAnswer(sql(s"SELECT '$c1' = '$c2' COLLATE SR_AI"), Row(true)) + } + // scalastyle:on nonascii + } + test("equality check respects collation") { Seq( ("utf8_binary", "aaa", "AAA", false), From 6393afada79ab4d6f7a45139017b21021ecfaec1 Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Tue, 17 Sep 2024 14:47:20 +0200 Subject: [PATCH 045/189] [SPARK-49505][SQL] Create new SQL functions "randstr" and "uniform" to generate random strings or numbers within ranges ### What changes were proposed in this pull request? This PR introduces two new SQL functions "randstr" and "uniform" to generate random strings or numbers within ranges. * The "randstr" function returns a string of the specified length whose characters are chosen uniformly at random from the following pool of characters: 0-9, a-z, A-Z. The random seed is optional. The string length must be a constant two-byte or four-byte integer (SMALLINT or INT, respectively). * The "uniform" function returns a random value with independent and identically distributed values with the specified range of numbers. The random seed is optional. The provided numbers specifying the minimum and maximum values of the range must be constant. If both of these numbers are integers, then the result will also be an integer. Otherwise if one or both of these are floating-point numbers, then the result will also be a floating-point number. For example: ``` SELECT randstr(5); > ceV0P SELECT randstr(10, 0) FROM VALUES (0), (1), (2) tab(col); > ceV0PXaR2I fYxVfArnv7 iSIv0VT2XL SELECT uniform(10, 20.0F); > 17.604954 SELECT uniform(10, 20, 0) FROM VALUES (0), (1), (2) tab(col); > 15 16 17 ``` ### Why are the changes needed? This improves the SQL functionality of Apache Spark and improves its parity with other systems: * https://clickhouse.com/docs/en/sql-reference/functions/random-functions#randuniform * https://docs.snowflake.com/en/sql-reference/functions/uniform * https://www.microfocus.com/documentation/silk-test/21.0.2/en/silktestclassic-help-en/STCLASSIC-8BFE8661-RANDSTRFUNCTION-REF.html * https://docs.snowflake.com/en/sql-reference/functions/randstr ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds golden file based test coverage. ### Was this patch authored or co-authored using generative AI tooling? Not this time. Closes #48004 from dtenedor/uniform-randstr-functions. Authored-by: Daniel Tenedorio Signed-off-by: Max Gekk --- .../catalyst/analysis/FunctionRegistry.scala | 2 + .../expressions/randomExpressions.scala | 264 +++++++++- .../catalyst/expressions/RandomSuite.scala | 24 + .../sql-functions/sql-expression-schema.md | 2 + .../sql-tests/analyzer-results/random.sql.out | 401 +++++++++++++++ .../resources/sql-tests/inputs/random.sql | 42 +- .../sql-tests/results/random.sql.out | 469 ++++++++++++++++++ 7 files changed, 1191 insertions(+), 13 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala index 5a3c4b0ec8696..d03d8114e9976 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala @@ -384,7 +384,9 @@ object FunctionRegistry { expression[Rand]("rand"), expression[Rand]("random", true, Some("3.0.0")), expression[Randn]("randn"), + expression[RandStr]("randstr"), expression[Stack]("stack"), + expression[Uniform]("uniform"), expression[ZeroIfNull]("zeroifnull"), CaseWhen.registryEntry, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala index f5db972a28643..ea9ca451c2cb1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala @@ -17,13 +17,18 @@ package org.apache.spark.sql.catalyst.expressions +import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.analysis.UnresolvedSeed +import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, UnresolvedSeed} +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.expressions.ExpectsInputTypes.{ordinalNumber, toSQLExpr, toSQLType} import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral} import org.apache.spark.sql.catalyst.expressions.codegen.Block._ -import org.apache.spark.sql.catalyst.trees.TreePattern.{EXPRESSION_WITH_RANDOM_SEED, TreePattern} +import org.apache.spark.sql.catalyst.trees.{BinaryLike, TernaryLike, UnaryLike} +import org.apache.spark.sql.catalyst.trees.TreePattern.{EXPRESSION_WITH_RANDOM_SEED, RUNTIME_REPLACEABLE, TreePattern} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.random.XORShiftRandom /** @@ -33,8 +38,7 @@ import org.apache.spark.util.random.XORShiftRandom * * Since this expression is stateful, it cannot be a case object. */ -abstract class RDG extends UnaryExpression with ExpectsInputTypes with Nondeterministic - with ExpressionWithRandomSeed { +trait RDG extends Expression with ExpressionWithRandomSeed { /** * Record ID within each partition. By being transient, the Random Number Generator is * reset every time we serialize and deserialize and initialize it. @@ -43,12 +47,6 @@ abstract class RDG extends UnaryExpression with ExpectsInputTypes with Nondeterm override def stateful: Boolean = true - override protected def initializeInternal(partitionIndex: Int): Unit = { - rng = new XORShiftRandom(seed + partitionIndex) - } - - override def seedExpression: Expression = child - @transient protected lazy val seed: Long = seedExpression match { case e if e.dataType == IntegerType => e.eval().asInstanceOf[Int] case e if e.dataType == LongType => e.eval().asInstanceOf[Long] @@ -57,6 +55,15 @@ abstract class RDG extends UnaryExpression with ExpectsInputTypes with Nondeterm override def nullable: Boolean = false override def dataType: DataType = DoubleType +} + +abstract class NondeterministicUnaryRDG + extends RDG with UnaryLike[Expression] with Nondeterministic with ExpectsInputTypes { + override def seedExpression: Expression = child + + override protected def initializeInternal(partitionIndex: Int): Unit = { + rng = new XORShiftRandom(seed + partitionIndex) + } override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(IntegerType, LongType)) } @@ -99,7 +106,7 @@ private[catalyst] object ExpressionWithRandomSeed { since = "1.5.0", group = "math_funcs") // scalastyle:on line.size.limit -case class Rand(child: Expression, hideSeed: Boolean = false) extends RDG { +case class Rand(child: Expression, hideSeed: Boolean = false) extends NondeterministicUnaryRDG { def this() = this(UnresolvedSeed, true) @@ -150,7 +157,7 @@ object Rand { since = "1.5.0", group = "math_funcs") // scalastyle:on line.size.limit -case class Randn(child: Expression, hideSeed: Boolean = false) extends RDG { +case class Randn(child: Expression, hideSeed: Boolean = false) extends NondeterministicUnaryRDG { def this() = this(UnresolvedSeed, true) @@ -181,3 +188,236 @@ case class Randn(child: Expression, hideSeed: Boolean = false) extends RDG { object Randn { def apply(seed: Long): Randn = Randn(Literal(seed, LongType)) } + +@ExpressionDescription( + usage = """ + _FUNC_(min, max[, seed]) - Returns a random value with independent and identically + distributed (i.i.d.) values with the specified range of numbers. The random seed is optional. + The provided numbers specifying the minimum and maximum values of the range must be constant. + If both of these numbers are integers, then the result will also be an integer. Otherwise if + one or both of these are floating-point numbers, then the result will also be a floating-point + number. + """, + examples = """ + Examples: + > SELECT _FUNC_(10, 20, 0) > 0 AS result; + true + """, + since = "4.0.0", + group = "math_funcs") +case class Uniform(min: Expression, max: Expression, seedExpression: Expression) + extends RuntimeReplaceable with TernaryLike[Expression] with RDG { + def this(min: Expression, max: Expression) = this(min, max, UnresolvedSeed) + + final override lazy val deterministic: Boolean = false + override val nodePatterns: Seq[TreePattern] = + Seq(RUNTIME_REPLACEABLE, EXPRESSION_WITH_RANDOM_SEED) + + override val dataType: DataType = { + val first = min.dataType + val second = max.dataType + (min.dataType, max.dataType) match { + case _ if !seedExpression.resolved || seedExpression.dataType == NullType => + NullType + case (_, NullType) | (NullType, _) => NullType + case (_, LongType) | (LongType, _) + if Seq(first, second).forall(integer) => LongType + case (_, IntegerType) | (IntegerType, _) + if Seq(first, second).forall(integer) => IntegerType + case (_, ShortType) | (ShortType, _) + if Seq(first, second).forall(integer) => ShortType + case (_, DoubleType) | (DoubleType, _) => DoubleType + case (_, FloatType) | (FloatType, _) => FloatType + case _ => + throw SparkException.internalError( + s"Unexpected argument data types: ${min.dataType}, ${max.dataType}") + } + } + + private def integer(t: DataType): Boolean = t match { + case _: ShortType | _: IntegerType | _: LongType => true + case _ => false + } + + override def checkInputDataTypes(): TypeCheckResult = { + var result: TypeCheckResult = TypeCheckResult.TypeCheckSuccess + def requiredType = "integer or floating-point" + Seq((min, "min", 0), + (max, "max", 1), + (seedExpression, "seed", 2)).foreach { + case (expr: Expression, name: String, index: Int) => + if (result == TypeCheckResult.TypeCheckSuccess) { + if (!expr.foldable) { + result = DataTypeMismatch( + errorSubClass = "NON_FOLDABLE_INPUT", + messageParameters = Map( + "inputName" -> name, + "inputType" -> requiredType, + "inputExpr" -> toSQLExpr(expr))) + } else expr.dataType match { + case _: ShortType | _: IntegerType | _: LongType | _: FloatType | _: DoubleType | + _: NullType => + case _ => + result = DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(index), + "requiredType" -> requiredType, + "inputSql" -> toSQLExpr(expr), + "inputType" -> toSQLType(expr.dataType))) + } + } + } + result + } + + override def first: Expression = min + override def second: Expression = max + override def third: Expression = seedExpression + + override def withNewSeed(newSeed: Long): Expression = + Uniform(min, max, Literal(newSeed, LongType)) + + override def withNewChildrenInternal( + newFirst: Expression, newSecond: Expression, newThird: Expression): Expression = + Uniform(newFirst, newSecond, newThird) + + override def replacement: Expression = { + if (Seq(min, max, seedExpression).exists(_.dataType == NullType)) { + Literal(null) + } else { + def cast(e: Expression, to: DataType): Expression = if (e.dataType == to) e else Cast(e, to) + cast(Add( + cast(min, DoubleType), + Multiply( + Subtract( + cast(max, DoubleType), + cast(min, DoubleType)), + Rand(seed))), + dataType) + } + } +} + +@ExpressionDescription( + usage = """ + _FUNC_(length[, seed]) - Returns a string of the specified length whose characters are chosen + uniformly at random from the following pool of characters: 0-9, a-z, A-Z. The random seed is + optional. The string length must be a constant two-byte or four-byte integer (SMALLINT or INT, + respectively). + """, + examples = + """ + Examples: + > SELECT _FUNC_(3, 0) AS result; + ceV + """, + since = "4.0.0", + group = "string_funcs") +case class RandStr(length: Expression, override val seedExpression: Expression) + extends ExpressionWithRandomSeed with BinaryLike[Expression] with Nondeterministic { + def this(length: Expression) = this(length, UnresolvedSeed) + + override def nullable: Boolean = false + override def dataType: DataType = StringType + override def stateful: Boolean = true + override def left: Expression = length + override def right: Expression = seedExpression + + /** + * Record ID within each partition. By being transient, the Random Number Generator is + * reset every time we serialize and deserialize and initialize it. + */ + @transient protected var rng: XORShiftRandom = _ + + @transient protected lazy val seed: Long = seedExpression match { + case e if e.dataType == IntegerType => e.eval().asInstanceOf[Int] + case e if e.dataType == LongType => e.eval().asInstanceOf[Long] + } + override protected def initializeInternal(partitionIndex: Int): Unit = { + rng = new XORShiftRandom(seed + partitionIndex) + } + + override def withNewSeed(newSeed: Long): Expression = RandStr(length, Literal(newSeed, LongType)) + override def withNewChildrenInternal(newFirst: Expression, newSecond: Expression): Expression = + RandStr(newFirst, newSecond) + + override def checkInputDataTypes(): TypeCheckResult = { + var result: TypeCheckResult = TypeCheckResult.TypeCheckSuccess + def requiredType = "INT or SMALLINT" + Seq((length, "length", 0), + (seedExpression, "seedExpression", 1)).foreach { + case (expr: Expression, name: String, index: Int) => + if (result == TypeCheckResult.TypeCheckSuccess) { + if (!expr.foldable) { + result = DataTypeMismatch( + errorSubClass = "NON_FOLDABLE_INPUT", + messageParameters = Map( + "inputName" -> name, + "inputType" -> requiredType, + "inputExpr" -> toSQLExpr(expr))) + } else expr.dataType match { + case _: ShortType | _: IntegerType => + case _: LongType if index == 1 => + case _ => + result = DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(index), + "requiredType" -> requiredType, + "inputSql" -> toSQLExpr(expr), + "inputType" -> toSQLType(expr.dataType))) + } + } + } + result + } + + override def evalInternal(input: InternalRow): Any = { + val numChars = length.eval(input).asInstanceOf[Number].intValue() + val bytes = new Array[Byte](numChars) + (0 until numChars).foreach { i => + // We generate a random number between 0 and 61, inclusive. Between the 62 different choices + // we choose 0-9, a-z, or A-Z, where each category comprises 10 choices, 26 choices, or 26 + // choices, respectively (10 + 26 + 26 = 62). + val num = (rng.nextInt() % 62).abs + num match { + case _ if num < 10 => + bytes.update(i, ('0' + num).toByte) + case _ if num < 36 => + bytes.update(i, ('a' + num - 10).toByte) + case _ => + bytes.update(i, ('A' + num - 36).toByte) + } + } + val result: UTF8String = UTF8String.fromBytes(bytes.toArray) + result + } + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val className = classOf[XORShiftRandom].getName + val rngTerm = ctx.addMutableState(className, "rng") + ctx.addPartitionInitializationStatement( + s"$rngTerm = new $className(${seed}L + partitionIndex);") + val eval = length.genCode(ctx) + ev.copy(code = + code""" + |${eval.code} + |int length = (int)(${eval.value}); + |char[] chars = new char[length]; + |for (int i = 0; i < length; i++) { + | int v = Math.abs($rngTerm.nextInt() % 62); + | if (v < 10) { + | chars[i] = (char)('0' + v); + | } else if (v < 36) { + | chars[i] = (char)('a' + (v - 10)); + | } else { + | chars[i] = (char)('A' + (v - 36)); + | } + |} + |UTF8String ${ev.value} = UTF8String.fromString(new String(chars)); + |boolean ${ev.isNull} = false; + |""".stripMargin, + isNull = FalseLiteral) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala index 2aa53f581555f..2d58d9d3136aa 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.catalyst.expressions import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.CatalystTypeConverters import org.apache.spark.sql.types.{IntegerType, LongType} class RandomSuite extends SparkFunSuite with ExpressionEvalHelper { @@ -41,4 +42,27 @@ class RandomSuite extends SparkFunSuite with ExpressionEvalHelper { assert(Rand(Literal(1L), false).sql === "rand(1L)") assert(Randn(Literal(1L), false).sql === "randn(1L)") } + + test("SPARK-49505: Test the RANDSTR and UNIFORM SQL functions without codegen") { + // Note that we use a seed of zero in these tests to keep the results deterministic. + def testRandStr(first: Any, result: Any): Unit = { + checkEvaluationWithoutCodegen( + RandStr(Literal(first), Literal(0)), CatalystTypeConverters.convertToCatalyst(result)) + } + testRandStr(1, "c") + testRandStr(5, "ceV0P") + testRandStr(10, "ceV0PXaR2I") + testRandStr(10L, "ceV0PXaR2I") + + def testUniform(first: Any, second: Any, result: Any): Unit = { + checkEvaluationWithoutCodegen( + Uniform(Literal(first), Literal(second), Literal(0)).replacement, + CatalystTypeConverters.convertToCatalyst(result)) + } + testUniform(0, 1, 0) + testUniform(0, 10, 7) + testUniform(0L, 10L, 7L) + testUniform(10.0F, 20.0F, 17.604954F) + testUniform(10L, 20.0F, 17.604954F) + } } diff --git a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md index f53b3874e6b8c..5ad1380e1fb82 100644 --- a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md +++ b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md @@ -265,6 +265,7 @@ | org.apache.spark.sql.catalyst.expressions.RaiseErrorExpressionBuilder | raise_error | SELECT raise_error('custom error message') | struct | | org.apache.spark.sql.catalyst.expressions.Rand | rand | SELECT rand() | struct | | org.apache.spark.sql.catalyst.expressions.Rand | random | SELECT random() | struct | +| org.apache.spark.sql.catalyst.expressions.RandStr | randstr | SELECT randstr(3, 0) AS result | struct | | org.apache.spark.sql.catalyst.expressions.Randn | randn | SELECT randn() | struct | | org.apache.spark.sql.catalyst.expressions.Rank | rank | SELECT a, b, rank(b) OVER (PARTITION BY a ORDER BY b) FROM VALUES ('A1', 2), ('A1', 1), ('A2', 3), ('A1', 1) tab(a, b) | struct | | org.apache.spark.sql.catalyst.expressions.RegExpCount | regexp_count | SELECT regexp_count('Steven Jones and Stephen Smith are the best players', 'Ste(v|ph)en') | struct | @@ -367,6 +368,7 @@ | org.apache.spark.sql.catalyst.expressions.UnaryMinus | negative | SELECT negative(1) | struct | | org.apache.spark.sql.catalyst.expressions.UnaryPositive | positive | SELECT positive(1) | struct<(+ 1):int> | | org.apache.spark.sql.catalyst.expressions.Unhex | unhex | SELECT decode(unhex('537061726B2053514C'), 'UTF-8') | struct | +| org.apache.spark.sql.catalyst.expressions.Uniform | uniform | SELECT uniform(10, 20, 0) > 0 AS result | struct | | org.apache.spark.sql.catalyst.expressions.UnixDate | unix_date | SELECT unix_date(DATE("1970-01-02")) | struct | | org.apache.spark.sql.catalyst.expressions.UnixMicros | unix_micros | SELECT unix_micros(TIMESTAMP('1970-01-01 00:00:01Z')) | struct | | org.apache.spark.sql.catalyst.expressions.UnixMillis | unix_millis | SELECT unix_millis(TIMESTAMP('1970-01-01 00:00:01Z')) | struct | diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out index 3cacbdc141053..133cd6a60a4fb 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out @@ -93,3 +93,404 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "fragment" : "rand('1')" } ] } + + +-- !query +SELECT uniform(0, 1, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0, 10, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0L, 10L, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0, 10L, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0, 10S, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10, 20, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10.0F, 20.0F, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10.0D, 20.0D, CAST(3 / 7 AS LONG)) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10, 20.0F, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10, 20, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10, 20.0F) IS NOT NULL AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(NULL, 1, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0, NULL, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(0, 1, NULL) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT uniform(10, 20, col) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "seed", + "inputType" : "integer or floating-point", + "sqlExpr" : "\"uniform(10, 20, col)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 27, + "fragment" : "uniform(10, 20, col)" + } ] +} + + +-- !query +SELECT uniform(col, 10, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "min", + "inputType" : "integer or floating-point", + "sqlExpr" : "\"uniform(col, 10, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 26, + "fragment" : "uniform(col, 10, 0)" + } ] +} + + +-- !query +SELECT uniform(10) AS result +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "1", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[2, 3]", + "functionName" : "`uniform`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 18, + "fragment" : "uniform(10)" + } ] +} + + +-- !query +SELECT uniform(10, 20, 30, 40) AS result +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "4", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[2, 3]", + "functionName" : "`uniform`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 30, + "fragment" : "uniform(10, 20, 30, 40)" + } ] +} + + +-- !query +SELECT randstr(1, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(5, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(10, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(10S, 0) AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(10, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(10) IS NOT NULL AS result +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT randstr(10L, 0) AS result +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10\"", + "inputType" : "\"BIGINT\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 22, + "fragment" : "randstr(10L, 0)" + } ] +} + + +-- !query +SELECT randstr(10.0F, 0) AS result +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10.0\"", + "inputType" : "\"FLOAT\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10.0, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10.0F, 0)" + } ] +} + + +-- !query +SELECT randstr(10.0D, 0) AS result +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10.0\"", + "inputType" : "\"DOUBLE\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10.0, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10.0D, 0)" + } ] +} + + +-- !query +SELECT randstr(NULL, 0) AS result +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"NULL\"", + "inputType" : "\"VOID\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(NULL, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(NULL, 0)" + } ] +} + + +-- !query +SELECT randstr(0, NULL) AS result +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"NULL\"", + "inputType" : "\"VOID\"", + "paramIndex" : "second", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(0, NULL)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(0, NULL)" + } ] +} + + +-- !query +SELECT randstr(col, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "length", + "inputType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(col, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 22, + "fragment" : "randstr(col, 0)" + } ] +} + + +-- !query +SELECT randstr(10, col) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "seedExpression", + "inputType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10, col)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(10, col)" + } ] +} + + +-- !query +SELECT randstr(10, 0, 1) AS result +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "3", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[1, 2]", + "functionName" : "`randstr`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10, 0, 1)" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/inputs/random.sql b/sql/core/src/test/resources/sql-tests/inputs/random.sql index a1aae7b8759dc..a71b0293295fc 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/random.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/random.sql @@ -14,4 +14,44 @@ SELECT randn(NULL); SELECT randn(cast(NULL AS long)); -- randn unsupported data type -SELECT rand('1') +SELECT rand('1'); + +-- The uniform random number generation function supports generating random numbers within a +-- specified range. We use a seed of zero for these queries to keep tests deterministic. +SELECT uniform(0, 1, 0) AS result; +SELECT uniform(0, 10, 0) AS result; +SELECT uniform(0L, 10L, 0) AS result; +SELECT uniform(0, 10L, 0) AS result; +SELECT uniform(0, 10S, 0) AS result; +SELECT uniform(10, 20, 0) AS result; +SELECT uniform(10.0F, 20.0F, 0) AS result; +SELECT uniform(10.0D, 20.0D, CAST(3 / 7 AS LONG)) AS result; +SELECT uniform(10, 20.0F, 0) AS result; +SELECT uniform(10, 20, 0) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT uniform(10, 20.0F) IS NOT NULL AS result; +-- Negative test cases for the uniform random number generator. +SELECT uniform(NULL, 1, 0) AS result; +SELECT uniform(0, NULL, 0) AS result; +SELECT uniform(0, 1, NULL) AS result; +SELECT uniform(10, 20, col) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT uniform(col, 10, 0) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT uniform(10) AS result; +SELECT uniform(10, 20, 30, 40) AS result; + +-- The randstr random string generation function supports generating random strings within a +-- specified length. We use a seed of zero for these queries to keep tests deterministic. +SELECT randstr(1, 0) AS result; +SELECT randstr(5, 0) AS result; +SELECT randstr(10, 0) AS result; +SELECT randstr(10S, 0) AS result; +SELECT randstr(10, 0) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT randstr(10) IS NOT NULL AS result; +-- Negative test cases for the randstr random number generator. +SELECT randstr(10L, 0) AS result; +SELECT randstr(10.0F, 0) AS result; +SELECT randstr(10.0D, 0) AS result; +SELECT randstr(NULL, 0) AS result; +SELECT randstr(0, NULL) AS result; +SELECT randstr(col, 0) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT randstr(10, col) AS result FROM VALUES (0), (1), (2) tab(col); +SELECT randstr(10, 0, 1) AS result; diff --git a/sql/core/src/test/resources/sql-tests/results/random.sql.out b/sql/core/src/test/resources/sql-tests/results/random.sql.out index 16984de3ff257..0b4e5e078ee15 100644 --- a/sql/core/src/test/resources/sql-tests/results/random.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/random.sql.out @@ -113,3 +113,472 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "fragment" : "rand('1')" } ] } + + +-- !query +SELECT uniform(0, 1, 0) AS result +-- !query schema +struct +-- !query output +0 + + +-- !query +SELECT uniform(0, 10, 0) AS result +-- !query schema +struct +-- !query output +7 + + +-- !query +SELECT uniform(0L, 10L, 0) AS result +-- !query schema +struct +-- !query output +7 + + +-- !query +SELECT uniform(0, 10L, 0) AS result +-- !query schema +struct +-- !query output +7 + + +-- !query +SELECT uniform(0, 10S, 0) AS result +-- !query schema +struct +-- !query output +7 + + +-- !query +SELECT uniform(10, 20, 0) AS result +-- !query schema +struct +-- !query output +17 + + +-- !query +SELECT uniform(10.0F, 20.0F, 0) AS result +-- !query schema +struct +-- !query output +17.604954 + + +-- !query +SELECT uniform(10.0D, 20.0D, CAST(3 / 7 AS LONG)) AS result +-- !query schema +struct +-- !query output +17.604953758285916 + + +-- !query +SELECT uniform(10, 20.0F, 0) AS result +-- !query schema +struct +-- !query output +17.604954 + + +-- !query +SELECT uniform(10, 20, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct +-- !query output +15 +16 +17 + + +-- !query +SELECT uniform(10, 20.0F) IS NOT NULL AS result +-- !query schema +struct +-- !query output +true + + +-- !query +SELECT uniform(NULL, 1, 0) AS result +-- !query schema +struct +-- !query output +NULL + + +-- !query +SELECT uniform(0, NULL, 0) AS result +-- !query schema +struct +-- !query output +NULL + + +-- !query +SELECT uniform(0, 1, NULL) AS result +-- !query schema +struct +-- !query output +NULL + + +-- !query +SELECT uniform(10, 20, col) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "seed", + "inputType" : "integer or floating-point", + "sqlExpr" : "\"uniform(10, 20, col)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 27, + "fragment" : "uniform(10, 20, col)" + } ] +} + + +-- !query +SELECT uniform(col, 10, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "min", + "inputType" : "integer or floating-point", + "sqlExpr" : "\"uniform(col, 10, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 26, + "fragment" : "uniform(col, 10, 0)" + } ] +} + + +-- !query +SELECT uniform(10) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "1", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[2, 3]", + "functionName" : "`uniform`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 18, + "fragment" : "uniform(10)" + } ] +} + + +-- !query +SELECT uniform(10, 20, 30, 40) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "4", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[2, 3]", + "functionName" : "`uniform`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 30, + "fragment" : "uniform(10, 20, 30, 40)" + } ] +} + + +-- !query +SELECT randstr(1, 0) AS result +-- !query schema +struct +-- !query output +c + + +-- !query +SELECT randstr(5, 0) AS result +-- !query schema +struct +-- !query output +ceV0P + + +-- !query +SELECT randstr(10, 0) AS result +-- !query schema +struct +-- !query output +ceV0PXaR2I + + +-- !query +SELECT randstr(10S, 0) AS result +-- !query schema +struct +-- !query output +ceV0PXaR2I + + +-- !query +SELECT randstr(10, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct +-- !query output +ceV0PXaR2I +fYxVfArnv7 +iSIv0VT2XL + + +-- !query +SELECT randstr(10) IS NOT NULL AS result +-- !query schema +struct +-- !query output +true + + +-- !query +SELECT randstr(10L, 0) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10\"", + "inputType" : "\"BIGINT\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 22, + "fragment" : "randstr(10L, 0)" + } ] +} + + +-- !query +SELECT randstr(10.0F, 0) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10.0\"", + "inputType" : "\"FLOAT\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10.0, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10.0F, 0)" + } ] +} + + +-- !query +SELECT randstr(10.0D, 0) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"10.0\"", + "inputType" : "\"DOUBLE\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10.0, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10.0D, 0)" + } ] +} + + +-- !query +SELECT randstr(NULL, 0) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"NULL\"", + "inputType" : "\"VOID\"", + "paramIndex" : "first", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(NULL, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(NULL, 0)" + } ] +} + + +-- !query +SELECT randstr(0, NULL) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"NULL\"", + "inputType" : "\"VOID\"", + "paramIndex" : "second", + "requiredType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(0, NULL)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(0, NULL)" + } ] +} + + +-- !query +SELECT randstr(col, 0) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "length", + "inputType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(col, 0)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 22, + "fragment" : "randstr(col, 0)" + } ] +} + + +-- !query +SELECT randstr(10, col) AS result FROM VALUES (0), (1), (2) tab(col) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + "sqlState" : "42K09", + "messageParameters" : { + "inputExpr" : "\"col\"", + "inputName" : "seedExpression", + "inputType" : "INT or SMALLINT", + "sqlExpr" : "\"randstr(10, col)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "randstr(10, col)" + } ] +} + + +-- !query +SELECT randstr(10, 0, 1) AS result +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "3", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "[1, 2]", + "functionName" : "`randstr`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 24, + "fragment" : "randstr(10, 0, 1)" + } ] +} From c38844c9ecc6dd648500b2ef6ff01acbe46255f4 Mon Sep 17 00:00:00 2001 From: Zhihong Yu Date: Tue, 17 Sep 2024 10:58:05 -0700 Subject: [PATCH 046/189] [SPARK-49687][SQL] Delay sorting in `validateAndMaybeEvolveStateSchema` ### What changes were proposed in this pull request? In `validateAndMaybeEvolveStateSchema`, existing schema and new schema are sorted by column family name. The sorting can be delayed until `createSchemaFile` is called. When computing `colFamiliesAddedOrRemoved`, we can use `toSet` to compare column families. ### Why are the changes needed? This would make `validateAndMaybeEvolveStateSchema` faster. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48116 from tedyu/ty-comp-chk. Authored-by: Zhihong Yu Signed-off-by: Dongjoon Hyun --- .../state/StateSchemaCompatibilityChecker.scala | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala index 3a1793f71794f..721d72b6a0991 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala @@ -168,12 +168,12 @@ class StateSchemaCompatibilityChecker( newStateSchema: List[StateStoreColFamilySchema], ignoreValueSchema: Boolean, stateSchemaVersion: Int): Boolean = { - val existingStateSchemaList = getExistingKeyAndValueSchema().sortBy(_.colFamilyName) - val newStateSchemaList = newStateSchema.sortBy(_.colFamilyName) + val existingStateSchemaList = getExistingKeyAndValueSchema() + val newStateSchemaList = newStateSchema if (existingStateSchemaList.isEmpty) { // write the schema file if it doesn't exist - createSchemaFile(newStateSchemaList, stateSchemaVersion) + createSchemaFile(newStateSchemaList.sortBy(_.colFamilyName), stateSchemaVersion) true } else { // validate if the new schema is compatible with the existing schema @@ -188,9 +188,9 @@ class StateSchemaCompatibilityChecker( } } val colFamiliesAddedOrRemoved = - newStateSchemaList.map(_.colFamilyName) != existingStateSchemaList.map(_.colFamilyName) + (newStateSchemaList.map(_.colFamilyName).toSet != existingSchemaMap.keySet) if (stateSchemaVersion == SCHEMA_FORMAT_V3 && colFamiliesAddedOrRemoved) { - createSchemaFile(newStateSchemaList, stateSchemaVersion) + createSchemaFile(newStateSchemaList.sortBy(_.colFamilyName), stateSchemaVersion) } // TODO: [SPARK-49535] Write Schema files after schema has changed for StateSchemaV3 colFamiliesAddedOrRemoved From 6fc176f4f34d73d6f6975836951562243343ba9a Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Tue, 17 Sep 2024 17:09:09 -0400 Subject: [PATCH 047/189] [SPARK-49413][CONNECT][SQL] Create a shared RuntimeConfig interface ### What changes were proposed in this pull request? This PR introduces a shared RuntimeConfig interface. ### Why are the changes needed? We are creating a shared Scala Spark SQL interface for Classic and Connect. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47980 from hvanhovell/SPARK-49413. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/SparkSession.scala | 14 +-- .../ConnectRuntimeConfig.scala} | 70 ++---------- .../kafka010/KafkaMicroBatchSourceSuite.scala | 4 +- project/MimaExcludes.scala | 4 + .../org/apache/spark/sql/RuntimeConfig.scala | 105 ++++++++++++++++++ .../apache/spark/sql/api/SparkSession.scala | 13 ++- .../execution/ExecuteGrpcResponseSender.scala | 8 +- .../sql/connect/service/SessionHolder.scala | 4 +- .../spark/sql/connect/utils/ErrorUtils.scala | 9 +- .../SparkConnectSessionHolderSuite.scala | 2 +- .../scala/org/apache/spark/sql/Dataset.scala | 4 +- .../org/apache/spark/sql/SparkSession.scala | 19 +--- .../spark/sql/artifact/ArtifactManager.scala | 2 +- .../spark/sql/execution/CacheManager.scala | 2 +- .../spark/sql/execution/SQLExecution.scala | 2 +- .../spark/sql/execution/command/ddl.scala | 2 +- .../spark/sql/execution/command/views.scala | 2 +- .../execution/datasources/DataSource.scala | 3 +- .../binaryfile/BinaryFileFormat.scala | 2 +- .../spark/sql/execution/datasources/ddl.scala | 2 +- .../v2/state/StateDataSource.scala | 8 +- .../execution/streaming/AsyncLogPurge.scala | 3 +- .../streaming/MicroBatchExecution.scala | 2 +- .../sql/execution/streaming/OffsetSeq.scala | 23 ++-- .../execution/streaming/StreamExecution.scala | 2 +- .../streaming/WatermarkTracker.scala | 3 +- .../RuntimeConfigImpl.scala} | 101 ++--------------- .../sql/streaming/StreamingQueryManager.scala | 2 +- .../spark/sql/FileBasedDataSourceSuite.scala | 2 +- .../org/apache/spark/sql/GenTPCDSData.scala | 2 +- .../org/apache/spark/sql/JoinSuite.scala | 4 +- .../apache/spark/sql/RuntimeConfigSuite.scala | 3 +- .../org/apache/spark/sql/SQLQuerySuite.scala | 7 +- .../spark/sql/SparkSessionBuilderSuite.scala | 20 ++-- .../sql/SparkSessionExtensionSuite.scala | 6 +- .../sql/StatisticsCollectionTestBase.scala | 4 +- .../sql/connector/DataSourceV2SQLSuite.scala | 8 +- .../CoalesceShufflePartitionsSuite.scala | 4 +- .../columnar/InMemoryColumnarQuerySuite.scala | 2 +- .../columnar/PartitionBatchPruningSuite.scala | 4 +- .../datasources/ReadSchemaSuite.scala | 10 +- .../binaryfile/BinaryFileFormatSuite.scala | 2 +- .../execution/datasources/orc/OrcTest.scala | 2 +- .../datasources/parquet/ParquetIOSuite.scala | 4 +- .../StateDataSourceChangeDataReadSuite.scala | 2 +- .../v2/state/StateDataSourceReadSuite.scala | 4 +- .../RocksDBStateStoreIntegrationSuite.scala | 2 +- .../streaming/state/RocksDBSuite.scala | 11 +- .../sql/expressions/ExpressionInfoSuite.scala | 2 +- .../spark/sql/internal/SQLConfSuite.scala | 88 +++++++-------- .../spark/sql/sources/BucketedReadSuite.scala | 2 +- .../sql/streaming/FileStreamSinkSuite.scala | 2 +- .../sql/streaming/FileStreamSourceSuite.scala | 4 +- .../FlatMapGroupsWithStateSuite.scala | 2 +- .../spark/sql/streaming/StreamTest.scala | 5 +- .../streaming/TriggerAvailableNowSuite.scala | 2 +- .../spark/sql/test/SharedSparkSession.scala | 2 + .../spark/sql/hive/HiveSharedStateSuite.scala | 2 +- .../spark/sql/hive/HiveSparkSubmitSuite.scala | 2 +- .../execution/HiveSerDeReadWriteSuite.scala | 2 +- .../sql/hive/execution/SQLQuerySuite.scala | 4 +- 61 files changed, 317 insertions(+), 322 deletions(-) rename connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/{RuntimeConfig.scala => internal/ConnectRuntimeConfig.scala} (68%) create mode 100644 sql/api/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala rename sql/core/src/main/scala/org/apache/spark/sql/{RuntimeConfig.scala => internal/RuntimeConfigImpl.scala} (51%) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 209ec88618c43..989a7e0c174c5 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -41,7 +41,7 @@ import org.apache.spark.sql.connect.client.{ClassFinder, CloseableIterator, Spar import org.apache.spark.sql.connect.client.SparkConnectClient.Configuration import org.apache.spark.sql.connect.client.arrow.ArrowSerializer import org.apache.spark.sql.functions.lit -import org.apache.spark.sql.internal.{CatalogImpl, SessionCleaner, SqlApiConf} +import org.apache.spark.sql.internal.{CatalogImpl, ConnectRuntimeConfig, SessionCleaner, SqlApiConf} import org.apache.spark.sql.internal.ColumnNodeToProtoConverter.{toExpr, toTypedExpr} import org.apache.spark.sql.streaming.DataStreamReader import org.apache.spark.sql.streaming.StreamingQueryManager @@ -88,16 +88,8 @@ class SparkSession private[sql] ( client.hijackServerSideSessionIdForTesting(suffix) } - /** - * Runtime configuration interface for Spark. - * - * This is the interface through which the user can get and set all Spark configurations that - * are relevant to Spark SQL. When getting the value of a config, his defaults to the value set - * in server, if any. - * - * @since 3.4.0 - */ - val conf: RuntimeConfig = new RuntimeConfig(client) + /** @inheritdoc */ + val conf: RuntimeConfig = new ConnectRuntimeConfig(client) /** @inheritdoc */ @transient diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/ConnectRuntimeConfig.scala similarity index 68% rename from connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala rename to connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/ConnectRuntimeConfig.scala index f77dd512ef257..7578e2424fb42 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/ConnectRuntimeConfig.scala @@ -14,10 +14,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.sql +package org.apache.spark.sql.internal import org.apache.spark.connect.proto.{ConfigRequest, ConfigResponse, KeyValue} import org.apache.spark.internal.Logging +import org.apache.spark.sql.RuntimeConfig import org.apache.spark.sql.connect.client.SparkConnectClient /** @@ -25,61 +26,31 @@ import org.apache.spark.sql.connect.client.SparkConnectClient * * @since 3.4.0 */ -class RuntimeConfig private[sql] (client: SparkConnectClient) extends Logging { +class ConnectRuntimeConfig private[sql] (client: SparkConnectClient) + extends RuntimeConfig + with Logging { - /** - * Sets the given Spark runtime configuration property. - * - * @since 3.4.0 - */ + /** @inheritdoc */ def set(key: String, value: String): Unit = { executeConfigRequest { builder => builder.getSetBuilder.addPairsBuilder().setKey(key).setValue(value) } } - /** - * Sets the given Spark runtime configuration property. - * - * @since 3.4.0 - */ - def set(key: String, value: Boolean): Unit = set(key, String.valueOf(value)) - - /** - * Sets the given Spark runtime configuration property. - * - * @since 3.4.0 - */ - def set(key: String, value: Long): Unit = set(key, String.valueOf(value)) - - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @throws java.util.NoSuchElementException - * if the key is not set and does not have a default value - * @since 3.4.0 - */ + /** @inheritdoc */ @throws[NoSuchElementException]("if the key is not set") def get(key: String): String = getOption(key).getOrElse { throw new NoSuchElementException(key) } - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @since 3.4.0 - */ + /** @inheritdoc */ def get(key: String, default: String): String = { executeConfigRequestSingleValue { builder => builder.getGetWithDefaultBuilder.addPairsBuilder().setKey(key).setValue(default) } } - /** - * Returns all properties set in this conf. - * - * @since 3.4.0 - */ + /** @inheritdoc */ def getAll: Map[String, String] = { val response = executeConfigRequest { builder => builder.getGetAllBuilder @@ -92,11 +63,7 @@ class RuntimeConfig private[sql] (client: SparkConnectClient) extends Logging { builder.result() } - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @since 3.4.0 - */ + /** @inheritdoc */ def getOption(key: String): Option[String] = { val pair = executeConfigRequestSinglePair { builder => builder.getGetOptionBuilder.addKeys(key) @@ -108,27 +75,14 @@ class RuntimeConfig private[sql] (client: SparkConnectClient) extends Logging { } } - /** - * Resets the configuration property for the given key. - * - * @since 3.4.0 - */ + /** @inheritdoc */ def unset(key: String): Unit = { executeConfigRequest { builder => builder.getUnsetBuilder.addKeys(key) } } - /** - * Indicates whether the configuration property with the given key is modifiable in the current - * session. - * - * @return - * `true` if the configuration property is modifiable. For static SQL, Spark Core, invalid - * (not existing) and other non-modifiable configuration properties, the returned value is - * `false`. - * @since 3.4.0 - */ + /** @inheritdoc */ def isModifiable(key: String): Boolean = { val modifiable = executeConfigRequestSingleValue { builder => builder.getIsModifiableBuilder.addKeys(key) diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala index 9ae6a9290f80a..1d119de43970f 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala @@ -1156,7 +1156,7 @@ abstract class KafkaMicroBatchSourceSuiteBase extends KafkaSourceSuiteBase with test("allow group.id prefix") { // Group ID prefix is only supported by consumer based offset reader - if (spark.conf.get(SQLConf.USE_DEPRECATED_KAFKA_OFFSET_FETCHING)) { + if (sqlConf.getConf(SQLConf.USE_DEPRECATED_KAFKA_OFFSET_FETCHING)) { testGroupId("groupIdPrefix", (expected, actual) => { assert(actual.exists(_.startsWith(expected)) && !actual.exists(_ === expected), "Valid consumer groups don't contain the expected group id - " + @@ -1167,7 +1167,7 @@ abstract class KafkaMicroBatchSourceSuiteBase extends KafkaSourceSuiteBase with test("allow group.id override") { // Group ID override is only supported by consumer based offset reader - if (spark.conf.get(SQLConf.USE_DEPRECATED_KAFKA_OFFSET_FETCHING)) { + if (sqlConf.getConf(SQLConf.USE_DEPRECATED_KAFKA_OFFSET_FETCHING)) { testGroupId("kafka.group.id", (expected, actual) => { assert(actual.exists(_ === expected), "Valid consumer groups don't " + s"contain the expected group id - Valid consumer groups: $actual / " + diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 6eee1e759e5ea..68433b501bcc4 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -160,6 +160,10 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.DataFrameWriterV2"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.WriteConfigMethods"), + // SPARK-49413: Create a shared RuntimeConfig interface. + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.RuntimeConfig"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.RuntimeConfig$"), + // SPARK-49287: Shared Streaming interfaces ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.scheduler.SparkListenerEvent"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.ForeachWriter"), diff --git a/sql/api/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala b/sql/api/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala new file mode 100644 index 0000000000000..23a2774ebc3a5 --- /dev/null +++ b/sql/api/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql + +import org.apache.spark.annotation.Stable + +/** + * Runtime configuration interface for Spark. To access this, use `SparkSession.conf`. + * + * Options set here are automatically propagated to the Hadoop configuration during I/O. + * + * @since 2.0.0 + */ +@Stable +abstract class RuntimeConfig { + + /** + * Sets the given Spark runtime configuration property. + * + * @since 2.0.0 + */ + def set(key: String, value: String): Unit + + /** + * Sets the given Spark runtime configuration property. + * + * @since 2.0.0 + */ + def set(key: String, value: Boolean): Unit = { + set(key, value.toString) + } + + /** + * Sets the given Spark runtime configuration property. + * + * @since 2.0.0 + */ + def set(key: String, value: Long): Unit = { + set(key, value.toString) + } + + /** + * Returns the value of Spark runtime configuration property for the given key. + * + * @throws java.util.NoSuchElementException + * if the key is not set and does not have a default value + * @since 2.0.0 + */ + @throws[NoSuchElementException]("if the key is not set") + def get(key: String): String + + /** + * Returns the value of Spark runtime configuration property for the given key. + * + * @since 2.0.0 + */ + def get(key: String, default: String): String + + /** + * Returns all properties set in this conf. + * + * @since 2.0.0 + */ + def getAll: Map[String, String] + + /** + * Returns the value of Spark runtime configuration property for the given key. + * + * @since 2.0.0 + */ + def getOption(key: String): Option[String] + + /** + * Resets the configuration property for the given key. + * + * @since 2.0.0 + */ + def unset(key: String): Unit + + /** + * Indicates whether the configuration property with the given key is modifiable in the current + * session. + * + * @return + * `true` if the configuration property is modifiable. For static SQL, Spark Core, invalid + * (not existing) and other non-modifiable configuration properties, the returned value is + * `false`. + * @since 2.4.0 + */ + def isModifiable(key: String): Boolean +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index cf502c746d24e..0580931620aaa 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -26,7 +26,7 @@ import _root_.java.net.URI import _root_.java.util import org.apache.spark.annotation.{DeveloperApi, Experimental} -import org.apache.spark.sql.{Encoder, Row} +import org.apache.spark.sql.{Encoder, Row, RuntimeConfig} import org.apache.spark.sql.types.StructType /** @@ -58,6 +58,17 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C */ def version: String + /** + * Runtime configuration interface for Spark. + * + * This is the interface through which the user can get and set all Spark and Hadoop + * configurations that are relevant to Spark SQL. When getting the value of a config, this + * defaults to the value set in the underlying `SparkContext`, if any. + * + * @since 2.0.0 + */ + def conf: RuntimeConfig + /** * A collection of methods for registering user-defined functions (UDF). * diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteGrpcResponseSender.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteGrpcResponseSender.scala index 3e360372d5600..051093fcad277 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteGrpcResponseSender.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteGrpcResponseSender.scala @@ -142,7 +142,9 @@ private[connect] class ExecuteGrpcResponseSender[T <: Message]( * client, but rather enqueued to in the response observer. */ private def enqueueProgressMessage(force: Boolean = false): Unit = { - if (executeHolder.sessionHolder.session.conf.get(CONNECT_PROGRESS_REPORT_INTERVAL) > 0) { + val progressReportInterval = executeHolder.sessionHolder.session.sessionState.conf + .getConf(CONNECT_PROGRESS_REPORT_INTERVAL) + if (progressReportInterval > 0) { SparkConnectService.executionListener.foreach { listener => // It is possible, that the tracker is no longer available and in this // case we simply ignore it and do not send any progress message. This avoids @@ -240,8 +242,8 @@ private[connect] class ExecuteGrpcResponseSender[T <: Message]( // monitor, and will notify upon state change. if (response.isEmpty) { // Wake up more frequently to send the progress updates. - val progressTimeout = - executeHolder.sessionHolder.session.conf.get(CONNECT_PROGRESS_REPORT_INTERVAL) + val progressTimeout = executeHolder.sessionHolder.session.sessionState.conf + .getConf(CONNECT_PROGRESS_REPORT_INTERVAL) // If the progress feature is disabled, wait for the deadline. val timeout = if (progressTimeout > 0) { progressTimeout diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala index 0cb820b39e875..e56d66da3050d 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala @@ -444,8 +444,8 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio */ private[connect] def usePlanCache(rel: proto.Relation, cachePlan: Boolean)( transform: proto.Relation => LogicalPlan): LogicalPlan = { - val planCacheEnabled = - Option(session).forall(_.conf.get(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED, true)) + val planCacheEnabled = Option(session) + .forall(_.sessionState.conf.getConf(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED, true)) // We only cache plans that have a plan ID. val hasPlanId = rel.hasCommon && rel.getCommon.hasPlanId diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala index 355048cf30363..f1636ed1ef092 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala @@ -205,7 +205,9 @@ private[connect] object ErrorUtils extends Logging { case _ => } - if (sessionHolderOpt.exists(_.session.conf.get(Connect.CONNECT_ENRICH_ERROR_ENABLED))) { + val enrichErrorEnabled = sessionHolderOpt.exists( + _.session.sessionState.conf.getConf(Connect.CONNECT_ENRICH_ERROR_ENABLED)) + if (enrichErrorEnabled) { // Generate a new unique key for this exception. val errorId = UUID.randomUUID().toString @@ -216,9 +218,10 @@ private[connect] object ErrorUtils extends Logging { } lazy val stackTrace = Option(ExceptionUtils.getStackTrace(st)) + val stackTraceEnabled = sessionHolderOpt.exists( + _.session.sessionState.conf.getConf(SQLConf.PYSPARK_JVM_STACKTRACE_ENABLED)) val withStackTrace = - if (sessionHolderOpt.exists( - _.session.conf.get(SQLConf.PYSPARK_JVM_STACKTRACE_ENABLED) && stackTrace.nonEmpty)) { + if (stackTraceEnabled && stackTrace.nonEmpty) { val maxSize = Math.min( SparkEnv.get.conf.get(Connect.CONNECT_JVM_STACK_TRACE_MAX_SIZE), maxMetadataSize) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala index beebe5d2e2dc1..ed2f60afb0096 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala @@ -399,7 +399,7 @@ class SparkConnectSessionHolderSuite extends SharedSparkSession { test("Test session plan cache - disabled") { val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) // Disable plan cache of the session - sessionHolder.session.conf.set(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED, false) + sessionHolder.session.conf.set(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED.key, false) val planner = new SparkConnectPlanner(sessionHolder) val query = buildRelation("select 1") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 0fab60a948423..6e5dcc24e29dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -243,7 +243,7 @@ class Dataset[T] private[sql]( @transient private[sql] val logicalPlan: LogicalPlan = { val plan = queryExecution.commandExecuted - if (sparkSession.conf.get(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED)) { + if (sparkSession.sessionState.conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED)) { val dsIds = plan.getTagValue(Dataset.DATASET_ID_TAG).getOrElse(new HashSet[Long]) dsIds.add(id) plan.setTagValue(Dataset.DATASET_ID_TAG, dsIds) @@ -772,7 +772,7 @@ class Dataset[T] private[sql]( private def addDataFrameIdToCol(expr: NamedExpression): NamedExpression = { val newExpr = expr transform { case a: AttributeReference - if sparkSession.conf.get(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED) => + if sparkSession.sessionState.conf.getConf(SQLConf.FAIL_AMBIGUOUS_SELF_JOIN_ENABLED) => val metadata = new MetadataBuilder() .withMetadata(a.metadata) .putLong(Dataset.DATASET_ID_KEY, id) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index a7fb71d95d147..5746b942341fc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -173,16 +173,8 @@ class SparkSession private( @transient val sqlContext: SQLContext = new SQLContext(this) - /** - * Runtime configuration interface for Spark. - * - * This is the interface through which the user can get and set all Spark and Hadoop - * configurations that are relevant to Spark SQL. When getting the value of a config, - * this defaults to the value set in the underlying `SparkContext`, if any. - * - * @since 2.0.0 - */ - @transient lazy val conf: RuntimeConfig = new RuntimeConfig(sessionState.conf) + /** @inheritdoc */ + @transient lazy val conf: RuntimeConfig = new RuntimeConfigImpl(sessionState.conf) /** * An interface to register custom [[org.apache.spark.sql.util.QueryExecutionListener]]s @@ -745,7 +737,8 @@ class SparkSession private( } private[sql] def leafNodeDefaultParallelism: Int = { - conf.get(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM).getOrElse(sparkContext.defaultParallelism) + sessionState.conf.getConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM) + .getOrElse(sparkContext.defaultParallelism) } private[sql] object Converter extends ColumnNodeToExpressionConverter with Serializable { @@ -1110,13 +1103,13 @@ object SparkSession extends Logging { private[sql] def getOrCloneSessionWithConfigsOff( session: SparkSession, configurations: Seq[ConfigEntry[Boolean]]): SparkSession = { - val configsEnabled = configurations.filter(session.conf.get[Boolean]) + val configsEnabled = configurations.filter(session.sessionState.conf.getConf[Boolean]) if (configsEnabled.isEmpty) { session } else { val newSession = session.cloneSession() configsEnabled.foreach(conf => { - newSession.conf.set(conf, false) + newSession.sessionState.conf.setConf(conf, false) }) newSession } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 4eb7d4fa17eea..1ee960622fc2a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -324,7 +324,7 @@ class ArtifactManager(session: SparkSession) extends Logging { val fs = destFSPath.getFileSystem(hadoopConf) if (fs.isInstanceOf[LocalFileSystem]) { val allowDestLocalConf = - session.conf.get(SQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL) + session.sessionState.conf.getConf(SQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL) .getOrElse( session.conf.get("spark.connect.copyFromLocalToFs.allowDestLocal").contains("true")) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala index aae424afcb8ac..1bf6f4e4d7d9f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala @@ -474,7 +474,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { // Bucketed scan only has one time overhead but can have multi-times benefits in cache, // so we always do bucketed scan in a cached plan. var disableConfigs = Seq(SQLConf.AUTO_BUCKETED_SCAN_ENABLED) - if (!session.conf.get(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING)) { + if (!session.sessionState.conf.getConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING)) { // Allowing changing cached plan output partitioning might lead to regression as it introduces // extra shuffle disableConfigs = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala index 58fff2d4a1a29..12ff649b621e3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala @@ -87,7 +87,7 @@ object SQLExecution extends Logging { executionIdToQueryExecution.put(executionId, queryExecution) val originalInterruptOnCancel = sc.getLocalProperty(SPARK_JOB_INTERRUPT_ON_CANCEL) if (originalInterruptOnCancel == null) { - val interruptOnCancel = sparkSession.conf.get(SQLConf.INTERRUPT_ON_CANCEL) + val interruptOnCancel = sparkSession.sessionState.conf.getConf(SQLConf.INTERRUPT_ON_CANCEL) sc.setInterruptOnCancel(interruptOnCancel) } try { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala index 3f221bfa53051..814e56b204f9e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala @@ -861,7 +861,7 @@ case class RepairTableCommand( // Hive metastore may not have enough memory to handle millions of partitions in single RPC, // we should split them into smaller batches. Since Hive client is not thread safe, we cannot // do this in parallel. - val batchSize = spark.conf.get(SQLConf.ADD_PARTITION_BATCH_SIZE) + val batchSize = spark.sessionState.conf.getConf(SQLConf.ADD_PARTITION_BATCH_SIZE) partitionSpecsAndLocs.iterator.grouped(batchSize).foreach { batch => val now = MILLISECONDS.toSeconds(System.currentTimeMillis()) val parts = batch.map { case (spec, location) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala index e1061a46db7b0..071e3826b20a0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala @@ -135,7 +135,7 @@ case class CreateViewCommand( referredTempFunctions) catalog.createTempView(name.table, tableDefinition, overrideIfExists = replace) } else if (viewType == GlobalTempView) { - val db = sparkSession.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val db = sparkSession.sessionState.conf.getConf(StaticSQLConf.GLOBAL_TEMP_DATABASE) val viewIdent = TableIdentifier(name.table, Option(db)) val aliasedPlan = aliasPlan(sparkSession, analyzedPlan) val tableDefinition = createTemporaryViewRelation( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index d88b5ee8877d7..968c204841e46 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -267,7 +267,8 @@ case class DataSource( checkAndGlobPathIfNecessary(checkEmptyGlobPath = false, checkFilesExist = false) createInMemoryFileIndex(globbedPaths) }) - val forceNullable = sparkSession.conf.get(SQLConf.FILE_SOURCE_SCHEMA_FORCE_NULLABLE) + val forceNullable = sparkSession.sessionState.conf + .getConf(SQLConf.FILE_SOURCE_SCHEMA_FORCE_NULLABLE) val sourceDataSchema = if (forceNullable) dataSchema.asNullable else dataSchema SourceInfo( s"FileSource[$path]", diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala index cbff526592f92..54c100282e2db 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala @@ -98,7 +98,7 @@ class BinaryFileFormat extends FileFormat with DataSourceRegister { val broadcastedHadoopConf = sparkSession.sparkContext.broadcast(new SerializableConfiguration(hadoopConf)) val filterFuncs = filters.flatMap(filter => createFilterFunction(filter)) - val maxLength = sparkSession.conf.get(SOURCES_BINARY_FILE_MAX_LENGTH) + val maxLength = sparkSession.sessionState.conf.getConf(SOURCES_BINARY_FILE_MAX_LENGTH) file: PartitionedFile => { val path = file.toPath diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ddl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ddl.scala index fc6cba786c4ed..d9367d92d462e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ddl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ddl.scala @@ -115,7 +115,7 @@ case class CreateTempViewUsing( }.logicalPlan if (global) { - val db = sparkSession.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val db = sparkSession.sessionState.conf.getConf(StaticSQLConf.GLOBAL_TEMP_DATABASE) val viewIdent = TableIdentifier(tableIdent.table, Option(db)) val viewDefinition = createTemporaryViewRelation( viewIdent, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala index 83399e2cac01b..50b90641d309b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala @@ -25,7 +25,7 @@ import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path import org.apache.spark.internal.Logging -import org.apache.spark.sql.{RuntimeConfig, SparkSession} +import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.DataSourceOptions import org.apache.spark.sql.connector.catalog.{Table, TableProvider} import org.apache.spark.sql.connector.expressions.Transform @@ -119,9 +119,9 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging throw StateDataSourceErrors.offsetMetadataLogUnavailable(batchId, checkpointLocation) ) - val clonedRuntimeConf = new RuntimeConfig(session.sessionState.conf.clone()) - OffsetSeqMetadata.setSessionConf(metadata, clonedRuntimeConf) - StateStoreConf(clonedRuntimeConf.sqlConf) + val clonedSqlConf = session.sessionState.conf.clone() + OffsetSeqMetadata.setSessionConf(metadata, clonedSqlConf) + StateStoreConf(clonedSqlConf) case _ => throw StateDataSourceErrors.offsetLogUnavailable(batchId, checkpointLocation) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/AsyncLogPurge.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/AsyncLogPurge.scala index 06fdc6c53bc4e..cb7e71bda84dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/AsyncLogPurge.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/AsyncLogPurge.scala @@ -49,7 +49,8 @@ trait AsyncLogPurge extends Logging { // which are written per run. protected def purgeStatefulMetadata(plan: SparkPlan): Unit - protected lazy val useAsyncPurge: Boolean = sparkSession.conf.get(SQLConf.ASYNC_LOG_PURGE) + protected lazy val useAsyncPurge: Boolean = sparkSession.sessionState.conf + .getConf(SQLConf.ASYNC_LOG_PURGE) protected def purgeAsync(batchId: Long): Unit = { if (purgeRunning.compareAndSet(false, true)) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala index 285494543533c..053aef6ced3a6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/MicroBatchExecution.scala @@ -477,7 +477,7 @@ class MicroBatchExecution( // update offset metadata nextOffsets.metadata.foreach { metadata => - OffsetSeqMetadata.setSessionConf(metadata, sparkSessionToRunBatches.conf) + OffsetSeqMetadata.setSessionConf(metadata, sparkSessionToRunBatches.sessionState.conf) execCtx.offsetSeqMetadata = OffsetSeqMetadata( metadata.batchWatermarkMs, metadata.batchTimestampMs, sparkSessionToRunBatches.conf) watermarkTracker = WatermarkTracker(sparkSessionToRunBatches.conf) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala index f0be33ad9a9d8..d5facc245e72f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala @@ -26,6 +26,7 @@ import org.apache.spark.io.CompressionCodec import org.apache.spark.sql.RuntimeConfig import org.apache.spark.sql.connector.read.streaming.{Offset => OffsetV2, SparkDataStream} import org.apache.spark.sql.execution.streaming.state.{FlatMapGroupsWithStateExecHelper, StreamingAggregationStateManager, SymmetricHashJoinStateManager} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf._ @@ -135,20 +136,21 @@ object OffsetSeqMetadata extends Logging { } /** Set the SparkSession configuration with the values in the metadata */ - def setSessionConf(metadata: OffsetSeqMetadata, sessionConf: RuntimeConfig): Unit = { + def setSessionConf(metadata: OffsetSeqMetadata, sessionConf: SQLConf): Unit = { + val configs = sessionConf.getAllConfs OffsetSeqMetadata.relevantSQLConfs.map(_.key).foreach { confKey => metadata.conf.get(confKey) match { case Some(valueInMetadata) => // Config value exists in the metadata, update the session config with this value - val optionalValueInSession = sessionConf.getOption(confKey) - if (optionalValueInSession.isDefined && optionalValueInSession.get != valueInMetadata) { + val optionalValueInSession = sessionConf.getConfString(confKey, null) + if (optionalValueInSession != null && optionalValueInSession != valueInMetadata) { logWarning(log"Updating the value of conf '${MDC(CONFIG, confKey)}' in current " + - log"session from '${MDC(OLD_VALUE, optionalValueInSession.get)}' " + + log"session from '${MDC(OLD_VALUE, optionalValueInSession)}' " + log"to '${MDC(NEW_VALUE, valueInMetadata)}'.") } - sessionConf.set(confKey, valueInMetadata) + sessionConf.setConfString(confKey, valueInMetadata) case None => // For backward compatibility, if a config was not recorded in the offset log, @@ -157,14 +159,17 @@ object OffsetSeqMetadata extends Logging { relevantSQLConfDefaultValues.get(confKey) match { case Some(defaultValue) => - sessionConf.set(confKey, defaultValue) + sessionConf.setConfString(confKey, defaultValue) logWarning(log"Conf '${MDC(CONFIG, confKey)}' was not found in the offset log, " + log"using default value '${MDC(DEFAULT_VALUE, defaultValue)}'") case None => - val valueStr = sessionConf.getOption(confKey).map { v => - s" Using existing session conf value '$v'." - }.getOrElse { " No value set in session conf." } + val value = sessionConf.getConfString(confKey, null) + val valueStr = if (value != null) { + s" Using existing session conf value '$value'." + } else { + " No value set in session conf." + } logWarning(log"Conf '${MDC(CONFIG, confKey)}' was not found in the offset log. " + log"${MDC(TIP, valueStr)}") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala index 4b1b9e02a242a..8f030884ad33b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala @@ -483,7 +483,7 @@ abstract class StreamExecution( @throws[TimeoutException] protected def interruptAndAwaitExecutionThreadTermination(): Unit = { val timeout = math.max( - sparkSession.conf.get(SQLConf.STREAMING_STOP_TIMEOUT), 0) + sparkSession.sessionState.conf.getConf(SQLConf.STREAMING_STOP_TIMEOUT), 0) queryExecutionThread.interrupt() queryExecutionThread.join(timeout) if (queryExecutionThread.isAlive) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/WatermarkTracker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/WatermarkTracker.scala index 54c47ec4e6ed8..3e6f122f463d3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/WatermarkTracker.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/WatermarkTracker.scala @@ -135,7 +135,8 @@ object WatermarkTracker { // saved in the checkpoint (e.g., old checkpoints), then the default `min` policy is enforced // through defaults specified in OffsetSeqMetadata.setSessionConf(). val policyName = conf.get( - SQLConf.STREAMING_MULTIPLE_WATERMARK_POLICY, MultipleWatermarkPolicy.DEFAULT_POLICY_NAME) + SQLConf.STREAMING_MULTIPLE_WATERMARK_POLICY.key, + MultipleWatermarkPolicy.DEFAULT_POLICY_NAME) new WatermarkTracker(MultipleWatermarkPolicy(policyName)) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala similarity index 51% rename from sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala rename to sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala index ed8cf4f121f03..ca439cdb89958 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala @@ -15,15 +15,15 @@ * limitations under the License. */ -package org.apache.spark.sql +package org.apache.spark.sql.internal import scala.jdk.CollectionConverters._ import org.apache.spark.SPARK_DOC_ROOT import org.apache.spark.annotation.Stable -import org.apache.spark.internal.config.{ConfigEntry, OptionalConfigEntry} +import org.apache.spark.internal.config.ConfigEntry +import org.apache.spark.sql.RuntimeConfig import org.apache.spark.sql.errors.QueryCompilationErrors -import org.apache.spark.sql.internal.SQLConf /** * Runtime configuration interface for Spark. To access this, use `SparkSession.conf`. @@ -33,89 +33,26 @@ import org.apache.spark.sql.internal.SQLConf * @since 2.0.0 */ @Stable -class RuntimeConfig private[sql](val sqlConf: SQLConf = new SQLConf) { +class RuntimeConfigImpl private[sql](val sqlConf: SQLConf = new SQLConf) extends RuntimeConfig { - /** - * Sets the given Spark runtime configuration property. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def set(key: String, value: String): Unit = { requireNonStaticConf(key) sqlConf.setConfString(key, value) } - /** - * Sets the given Spark runtime configuration property. - * - * @since 2.0.0 - */ - def set(key: String, value: Boolean): Unit = { - set(key, value.toString) - } - - /** - * Sets the given Spark runtime configuration property. - * - * @since 2.0.0 - */ - def set(key: String, value: Long): Unit = { - set(key, value.toString) - } - - /** - * Sets the given Spark runtime configuration property. - */ - private[sql] def set[T](entry: ConfigEntry[T], value: T): Unit = { - requireNonStaticConf(entry.key) - sqlConf.setConf(entry, value) - } - - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @throws java.util.NoSuchElementException if the key is not set and does not have a default - * value - * @since 2.0.0 - */ + /** @inheritdoc */ @throws[NoSuchElementException]("if the key is not set") def get(key: String): String = { sqlConf.getConfString(key) } - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def get(key: String, default: String): String = { sqlConf.getConfString(key, default) } - /** - * Returns the value of Spark runtime configuration property for the given key. - */ - @throws[NoSuchElementException]("if the key is not set") - private[sql] def get[T](entry: ConfigEntry[T]): T = { - sqlConf.getConf(entry) - } - - private[sql] def get[T](entry: OptionalConfigEntry[T]): Option[T] = { - sqlConf.getConf(entry) - } - - /** - * Returns the value of Spark runtime configuration property for the given key. - */ - private[sql] def get[T](entry: ConfigEntry[T], default: T): T = { - sqlConf.getConf(entry, default) - } - - /** - * Returns all properties set in this conf. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def getAll: Map[String, String] = { sqlConf.getAllConfs } @@ -124,36 +61,20 @@ class RuntimeConfig private[sql](val sqlConf: SQLConf = new SQLConf) { getAll.asJava } - /** - * Returns the value of Spark runtime configuration property for the given key. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def getOption(key: String): Option[String] = { try Option(get(key)) catch { case _: NoSuchElementException => None } } - /** - * Resets the configuration property for the given key. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def unset(key: String): Unit = { requireNonStaticConf(key) sqlConf.unsetConf(key) } - /** - * Indicates whether the configuration property with the given key - * is modifiable in the current session. - * - * @return `true` if the configuration property is modifiable. For static SQL, Spark Core, - * invalid (not existing) and other non-modifiable configuration properties, - * the returned value is `false`. - * @since 2.4.0 - */ + /** @inheritdoc */ def isModifiable(key: String): Boolean = sqlConf.isModifiable(key) /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala index 55d2e639a56b1..3ab6d02f6b515 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala @@ -364,7 +364,7 @@ class StreamingQueryManager private[sql] ( .orElse(activeQueries.get(query.id)) // shouldn't be needed but paranoia ... val shouldStopActiveRun = - sparkSession.conf.get(SQLConf.STREAMING_STOP_ACTIVE_RUN_ON_RESTART) + sparkSession.sessionState.conf.getConf(SQLConf.STREAMING_STOP_ACTIVE_RUN_ON_RESTART) if (activeOption.isDefined) { if (shouldStopActiveRun) { val oldQuery = activeOption.get diff --git a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala index 2fe6a83427bca..e44bd5de4f4c4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala @@ -52,7 +52,7 @@ class FileBasedDataSourceSuite extends QueryTest override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(SQLConf.ORC_IMPLEMENTATION, "native") + spark.conf.set(SQLConf.ORC_IMPLEMENTATION.key, "native") } override def afterAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/GenTPCDSData.scala b/sql/core/src/test/scala/org/apache/spark/sql/GenTPCDSData.scala index 48a16f01d5749..6cd8ade41da14 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/GenTPCDSData.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/GenTPCDSData.scala @@ -225,7 +225,7 @@ class TPCDSTables(spark: SparkSession, dsdgenDir: String, scaleFactor: Int) // datagen speed files will be truncated to maxRecordsPerFile value, so the final // result will be the same. val numRows = data.count() - val maxRecordPerFile = spark.conf.get(SQLConf.MAX_RECORDS_PER_FILE) + val maxRecordPerFile = spark.sessionState.conf.getConf(SQLConf.MAX_RECORDS_PER_FILE) if (maxRecordPerFile > 0 && numRows > maxRecordPerFile) { val numFiles = (numRows.toDouble/maxRecordPerFile).ceil.toInt diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala index fcb937d82ba42..0f5582def82da 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala @@ -597,10 +597,10 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan SQLConf.CROSS_JOINS_ENABLED.key -> "true") { assert(statisticSizeInByte(spark.table("testData2")) > - spark.conf.get[Long](SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) + sqlConf.getConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) assert(statisticSizeInByte(spark.table("testData")) < - spark.conf.get[Long](SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) + sqlConf.getConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) Seq( ("SELECT * FROM testData LEFT SEMI JOIN testData2 ON key = a", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/RuntimeConfigSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/RuntimeConfigSuite.scala index 4052130720811..352197f96acb6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/RuntimeConfigSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/RuntimeConfigSuite.scala @@ -19,12 +19,13 @@ package org.apache.spark.sql import org.apache.spark.SparkFunSuite import org.apache.spark.internal.config +import org.apache.spark.sql.internal.RuntimeConfigImpl import org.apache.spark.sql.internal.SQLConf.CHECKPOINT_LOCATION import org.apache.spark.sql.internal.StaticSQLConf.GLOBAL_TEMP_DATABASE class RuntimeConfigSuite extends SparkFunSuite { - private def newConf(): RuntimeConfig = new RuntimeConfig + private def newConf(): RuntimeConfig = new RuntimeConfigImpl() test("set and get") { val conf = newConf() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index 9beceda263797..ce88f7dc475d6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -2568,20 +2568,21 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark Seq(1 -> "a").toDF("i", "j").write.parquet(path.getCanonicalPath) val newSession = spark.newSession() + val newSqlConf = newSession.sessionState.conf val originalValue = newSession.sessionState.conf.runSQLonFile try { - newSession.conf.set(SQLConf.RUN_SQL_ON_FILES, false) + newSqlConf.setConf(SQLConf.RUN_SQL_ON_FILES, false) intercept[AnalysisException] { newSession.sql(s"SELECT i, j FROM parquet.`${path.getCanonicalPath}`") } - newSession.conf.set(SQLConf.RUN_SQL_ON_FILES, true) + newSqlConf.setConf(SQLConf.RUN_SQL_ON_FILES, true) checkAnswer( newSession.sql(s"SELECT i, j FROM parquet.`${path.getCanonicalPath}`"), Row(1, "a")) } finally { - newSession.conf.set(SQLConf.RUN_SQL_ON_FILES, originalValue) + newSqlConf.setConf(SQLConf.RUN_SQL_ON_FILES, originalValue) } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderSuite.scala index 4ac05373e5a34..d3117ec411feb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderSuite.scala @@ -201,10 +201,10 @@ class SparkSessionBuilderSuite extends SparkFunSuite with Eventually { .getOrCreate() assert(session.conf.get("spark.app.name") === "test-app-SPARK-31234") - assert(session.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31234") + assert(session.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31234") session.sql("RESET") assert(session.conf.get("spark.app.name") === "test-app-SPARK-31234") - assert(session.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31234") + assert(session.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31234") } test("SPARK-31354: SparkContext only register one SparkSession ApplicationEnd listener") { @@ -244,8 +244,8 @@ class SparkSessionBuilderSuite extends SparkFunSuite with Eventually { .builder() .config(GLOBAL_TEMP_DATABASE.key, "globalTempDB-SPARK-31532-1") .getOrCreate() - assert(session.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31532") - assert(session1.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31532") + assert(session.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31532") + assert(session1.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31532") // do not propagate static sql configs to the existing default session SparkSession.clearActiveSession() @@ -255,9 +255,9 @@ class SparkSessionBuilderSuite extends SparkFunSuite with Eventually { .config(GLOBAL_TEMP_DATABASE.key, value = "globalTempDB-SPARK-31532-2") .getOrCreate() - assert(!session.conf.get(WAREHOUSE_PATH).contains("SPARK-31532-db")) - assert(session.conf.get(WAREHOUSE_PATH) === session2.conf.get(WAREHOUSE_PATH)) - assert(session2.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31532") + assert(!session.conf.get(WAREHOUSE_PATH.key).contains("SPARK-31532-db")) + assert(session.conf.get(WAREHOUSE_PATH.key) === session2.conf.get(WAREHOUSE_PATH.key)) + assert(session2.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31532") } test("SPARK-31532: propagate static sql configs if no existing SparkSession") { @@ -275,8 +275,8 @@ class SparkSessionBuilderSuite extends SparkFunSuite with Eventually { .config(WAREHOUSE_PATH.key, "SPARK-31532-db-2") .getOrCreate() assert(session.conf.get("spark.app.name") === "test-app-SPARK-31532-2") - assert(session.conf.get(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31532-2") - assert(session.conf.get(WAREHOUSE_PATH) contains "SPARK-31532-db-2") + assert(session.conf.get(GLOBAL_TEMP_DATABASE.key) === "globalTempDB-SPARK-31532-2") + assert(session.conf.get(WAREHOUSE_PATH.key) contains "SPARK-31532-db-2") } test("SPARK-32062: reset listenerRegistered in SparkSession") { @@ -461,7 +461,7 @@ class SparkSessionBuilderSuite extends SparkFunSuite with Eventually { val expected = path.getFileSystem(hadoopConf).makeQualified(path).toString // session related configs assert(hadoopConf.get("hive.metastore.warehouse.dir") === expected) - assert(session.conf.get(WAREHOUSE_PATH) === expected) + assert(session.conf.get(WAREHOUSE_PATH.key) === expected) assert(session.sessionState.conf.warehousePath === expected) // shared configs diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index 322210bf5b59f..ba87028a71477 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -178,7 +178,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt MyColumnarRule(MyNewQueryStageRule(), MyNewQueryStageRule())) } withSession(extensions) { session => - session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true) + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) assert(session.sessionState.adaptiveRulesHolder.queryStagePrepRules .contains(MyQueryStagePrepRule())) assert(session.sessionState.columnarRules.contains( @@ -221,7 +221,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) } withSession(extensions) { session => - session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true) + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) session.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") assert(session.sessionState.columnarRules.contains( MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))) @@ -280,7 +280,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) } withSession(extensions) { session => - session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, enableAQE) + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, enableAQE) assert(session.sessionState.columnarRules.contains( MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule()))) import session.implicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala index ef8b66566f246..7fa29dd38fd96 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala @@ -366,7 +366,7 @@ abstract class StatisticsCollectionTestBase extends QueryTest with SQLTestUtils val stats = spark.table("ds_tbl").queryExecution.optimizedPlan.stats assert(stats.sizeInBytes > 0, "non-empty partitioned table should not report zero size.") - if (spark.conf.get(StaticSQLConf.CATALOG_IMPLEMENTATION) == "hive") { + if (spark.conf.get(StaticSQLConf.CATALOG_IMPLEMENTATION.key) == "hive") { sql("CREATE TABLE hive_tbl(i int) PARTITIONED BY (j int)") sql("INSERT INTO hive_tbl PARTITION(j=1) SELECT 1") val stats2 = spark.table("hive_tbl").queryExecution.optimizedPlan.stats @@ -381,7 +381,7 @@ abstract class StatisticsCollectionTestBase extends QueryTest with SQLTestUtils // Test data source table checkStatsConversion(tableName = "ds_tbl", isDatasourceTable = true) // Test hive serde table - if (spark.conf.get(StaticSQLConf.CATALOG_IMPLEMENTATION) == "hive") { + if (spark.conf.get(StaticSQLConf.CATALOG_IMPLEMENTATION.key) == "hive") { checkStatsConversion(tableName = "hive_tbl", isDatasourceTable = false) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index 5df7b62cfb285..7aaec6d500ba0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -2481,7 +2481,7 @@ class DataSourceV2SQLSuiteV1Filter } test("global temp view should not be masked by v2 catalog") { - val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE.key) registerCatalog(globalTempDB, classOf[InMemoryTableCatalog]) try { @@ -2495,7 +2495,7 @@ class DataSourceV2SQLSuiteV1Filter } test("SPARK-30104: global temp db is used as a table name under v2 catalog") { - val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE.key) val t = s"testcat.$globalTempDB" withTable(t) { sql(s"CREATE TABLE $t (id bigint, data string) USING foo") @@ -2506,7 +2506,7 @@ class DataSourceV2SQLSuiteV1Filter } test("SPARK-30104: v2 catalog named global_temp will be masked") { - val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE.key) registerCatalog(globalTempDB, classOf[InMemoryTableCatalog]) checkError( exception = intercept[AnalysisException] { @@ -2712,7 +2712,7 @@ class DataSourceV2SQLSuiteV1Filter parameters = Map("relationName" -> "`testcat`.`abc`"), context = ExpectedContext(fragment = "testcat.abc", start = 17, stop = 27)) - val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE) + val globalTempDB = spark.conf.get(StaticSQLConf.GLOBAL_TEMP_DATABASE.key) registerCatalog(globalTempDB, classOf[InMemoryTableCatalog]) withTempView("v") { sql("create global temp view v as select 1") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala index dc72b4a092aef..9ed4f1a006b2b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -317,7 +317,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper { import spark.implicits._ spark.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "1KB") spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key, "10KB") - spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR, 2.0) + spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR.key, "2.0") val df00 = spark.range(0, 1000, 2) .selectExpr("id as key", "id as value") .union(Seq.fill(100000)((600, 600)).toDF("key", "value")) @@ -345,7 +345,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper { import spark.implicits._ spark.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key, "100B") - spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR, 2.0) + spark.conf.set(SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR.key, "2.0") val df00 = spark.range(0, 10, 2) .selectExpr("id as key", "id as value") .union(Seq.fill(1000)((600, 600)).toDF("key", "value")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index ad755bf22ab09..0ba55382cd9a7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -150,7 +150,7 @@ class InMemoryColumnarQuerySuite extends QueryTest spark.catalog.cacheTable("sizeTst") assert( spark.table("sizeTst").queryExecution.analyzed.stats.sizeInBytes > - spark.conf.get(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) + sqlConf.getConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala index 885286843a143..88ff51d0ff4cf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala @@ -27,9 +27,9 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl import testImplicits._ - private lazy val originalColumnBatchSize = spark.conf.get(SQLConf.COLUMN_BATCH_SIZE) + private lazy val originalColumnBatchSize = spark.conf.get(SQLConf.COLUMN_BATCH_SIZE.key) private lazy val originalInMemoryPartitionPruning = - spark.conf.get(SQLConf.IN_MEMORY_PARTITION_PRUNING) + spark.conf.get(SQLConf.IN_MEMORY_PARTITION_PRUNING.key) private val testArrayData = (1 to 100).map { key => Tuple1(Array.fill(key)(key)) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ReadSchemaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ReadSchemaSuite.scala index fefb16a351fdb..c798196c4f0ee 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ReadSchemaSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ReadSchemaSuite.scala @@ -101,7 +101,7 @@ class OrcReadSchemaSuite override def beforeAll(): Unit = { super.beforeAll() - originalConf = spark.conf.get(SQLConf.ORC_VECTORIZED_READER_ENABLED) + originalConf = sqlConf.getConf(SQLConf.ORC_VECTORIZED_READER_ENABLED) spark.conf.set(SQLConf.ORC_VECTORIZED_READER_ENABLED.key, "false") } @@ -126,7 +126,7 @@ class VectorizedOrcReadSchemaSuite override def beforeAll(): Unit = { super.beforeAll() - originalConf = spark.conf.get(SQLConf.ORC_VECTORIZED_READER_ENABLED) + originalConf = sqlConf.getConf(SQLConf.ORC_VECTORIZED_READER_ENABLED) spark.conf.set(SQLConf.ORC_VECTORIZED_READER_ENABLED.key, "true") } @@ -169,7 +169,7 @@ class ParquetReadSchemaSuite override def beforeAll(): Unit = { super.beforeAll() - originalConf = spark.conf.get(SQLConf.PARQUET_VECTORIZED_READER_ENABLED) + originalConf = sqlConf.getConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED) spark.conf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, "false") } @@ -193,7 +193,7 @@ class VectorizedParquetReadSchemaSuite override def beforeAll(): Unit = { super.beforeAll() - originalConf = spark.conf.get(SQLConf.PARQUET_VECTORIZED_READER_ENABLED) + originalConf = sqlConf.getConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED) spark.conf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, "true") } @@ -217,7 +217,7 @@ class MergedParquetReadSchemaSuite override def beforeAll(): Unit = { super.beforeAll() - originalConf = spark.conf.get(SQLConf.PARQUET_SCHEMA_MERGING_ENABLED) + originalConf = sqlConf.getConf(SQLConf.PARQUET_SCHEMA_MERGING_ENABLED) spark.conf.set(SQLConf.PARQUET_SCHEMA_MERGING_ENABLED.key, "true") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala index 3dec1b9ff5cf2..deb62eb3ac234 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala @@ -346,7 +346,7 @@ class BinaryFileFormatSuite extends QueryTest with SharedSparkSession { } test("fail fast and do not attempt to read if a file is too big") { - assert(spark.conf.get(SOURCES_BINARY_FILE_MAX_LENGTH) === Int.MaxValue) + assert(sqlConf.getConf(SOURCES_BINARY_FILE_MAX_LENGTH) === Int.MaxValue) withTempPath { file => val path = file.getPath val content = "123".getBytes diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcTest.scala index 48b4f8d4bc015..b8669ee4d1ef1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcTest.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcTest.scala @@ -63,7 +63,7 @@ trait OrcTest extends QueryTest with FileBasedDataSourceTest with BeforeAndAfter protected override def beforeAll(): Unit = { super.beforeAll() - originalConfORCImplementation = spark.conf.get(ORC_IMPLEMENTATION) + originalConfORCImplementation = spark.sessionState.conf.getConf(ORC_IMPLEMENTATION) spark.conf.set(ORC_IMPLEMENTATION.key, orcImp) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala index 6c2f5a2d134db..0afa545595c77 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala @@ -846,7 +846,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession def checkCompressionCodec(codec: ParquetCompressionCodec): Unit = { withSQLConf(SQLConf.PARQUET_COMPRESSION.key -> codec.name()) { withParquetFile(data) { path => - assertResult(spark.conf.get(SQLConf.PARQUET_COMPRESSION).toUpperCase(Locale.ROOT)) { + assertResult(spark.conf.get(SQLConf.PARQUET_COMPRESSION.key).toUpperCase(Locale.ROOT)) { compressionCodecFor(path, codec.name()) } } @@ -855,7 +855,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession // Checks default compression codec checkCompressionCodec( - ParquetCompressionCodec.fromString(spark.conf.get(SQLConf.PARQUET_COMPRESSION))) + ParquetCompressionCodec.fromString(spark.conf.get(SQLConf.PARQUET_COMPRESSION.key))) ParquetCompressionCodec.availableCodecs.asScala.foreach(checkCompressionCodec(_)) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceChangeDataReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceChangeDataReadSuite.scala index 2858d356d4c9a..4833b8630134c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceChangeDataReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceChangeDataReadSuite.scala @@ -58,7 +58,7 @@ abstract class StateDataSourceChangeDataReaderSuite extends StateDataSourceTestB override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(SQLConf.STREAMING_NO_DATA_MICRO_BATCHES_ENABLED, false) + spark.conf.set(SQLConf.STREAMING_NO_DATA_MICRO_BATCHES_ENABLED.key, false) spark.conf.set(SQLConf.STATE_STORE_PROVIDER_CLASS.key, newStateStoreProvider().getClass.getName) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala index 97c88037a7171..af07707569500 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala @@ -942,7 +942,7 @@ abstract class StateDataSourceReadSuite extends StateDataSourceTestBase with Ass // skip version and operator ID to test out functionalities .load() - val numShufflePartitions = spark.conf.get(SQLConf.SHUFFLE_PARTITIONS) + val numShufflePartitions = sqlConf.getConf(SQLConf.SHUFFLE_PARTITIONS) val resultDf = stateReadDf .selectExpr("key.value AS key_value", "value.count AS value_count", "partition_id") @@ -966,7 +966,7 @@ abstract class StateDataSourceReadSuite extends StateDataSourceTestBase with Ass } test("partition_id column with stream-stream join") { - val numShufflePartitions = spark.conf.get(SQLConf.SHUFFLE_PARTITIONS) + val numShufflePartitions = sqlConf.getConf(SQLConf.SHUFFLE_PARTITIONS) withTempDir { tempDir => runStreamStreamJoinQueryWithOneThousandInputs(tempDir.getAbsolutePath) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala index 8fcd6edf1abb7..d20cfb04f8e81 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala @@ -119,7 +119,7 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest private def getFormatVersion(query: StreamingQuery): Int = { query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution.sparkSession - .conf.get(SQLConf.STATE_STORE_ROCKSDB_FORMAT_VERSION) + .sessionState.conf.getConf(SQLConf.STATE_STORE_ROCKSDB_FORMAT_VERSION) } testWithColumnFamilies("SPARK-36519: store RocksDB format version in the checkpoint", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala index 7ac574db98d45..691f18451af22 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala @@ -34,7 +34,7 @@ import org.rocksdb.CompressionType import org.scalactic.source.Position import org.scalatest.Tag -import org.apache.spark.SparkException +import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.sql.catalyst.util.quietly import org.apache.spark.sql.execution.streaming.{CreateAtomicTestManager, FileSystemBasedCheckpointFileManager} import org.apache.spark.sql.execution.streaming.CheckpointFileManager.{CancellableFSDataOutputStream, RenameBasedFSDataOutputStream} @@ -167,7 +167,10 @@ trait AlsoTestWithChangelogCheckpointingEnabled @SlowSQLTest class RocksDBSuite extends AlsoTestWithChangelogCheckpointingEnabled with SharedSparkSession { - sqlConf.setConf(SQLConf.STATE_STORE_PROVIDER_CLASS, classOf[RocksDBStateStoreProvider].getName) + override protected def sparkConf: SparkConf = { + super.sparkConf + .set(SQLConf.STATE_STORE_PROVIDER_CLASS, classOf[RocksDBStateStoreProvider].getName) + } testWithColumnFamilies( "RocksDB: check changelog and snapshot version", @@ -2157,9 +2160,7 @@ class RocksDBSuite extends AlsoTestWithChangelogCheckpointingEnabled with Shared } } - private def sqlConf = SQLConf.get.clone() - - private def dbConf = RocksDBConf(StateStoreConf(sqlConf)) + private def dbConf = RocksDBConf(StateStoreConf(SQLConf.get.clone())) def withDB[T]( remoteDir: String, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala index 898aeec22ad17..6eff610433c9c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala @@ -243,7 +243,7 @@ class ExpressionInfoSuite extends SparkFunSuite with SharedSparkSession { // Examples can change settings. We clone the session to prevent tests clashing. val clonedSpark = spark.cloneSession() // Coalescing partitions can change result order, so disable it. - clonedSpark.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED, false) + clonedSpark.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, false) val info = clonedSpark.sessionState.catalog.lookupFunctionInfo(funcId) val className = info.getClassName if (!ignoreSet.contains(className)) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala index d0d4dc6b344fc..82795e551b6bf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala @@ -47,7 +47,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { // Set a conf first. spark.conf.set(testKey, testVal) // Clear the conf. - spark.sessionState.conf.clear() + sqlConf.clear() // After clear, only overrideConfs used by unit test should be in the SQLConf. assert(spark.conf.getAll === TestSQLContext.overrideConfs) @@ -62,11 +62,11 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { assert(spark.conf.get(testKey, testVal + "_") === testVal) assert(spark.conf.getAll.contains(testKey)) - spark.sessionState.conf.clear() + sqlConf.clear() } test("parse SQL set commands") { - spark.sessionState.conf.clear() + sqlConf.clear() sql(s"set $testKey=$testVal") assert(spark.conf.get(testKey, testVal + "_") === testVal) assert(spark.conf.get(testKey, testVal + "_") === testVal) @@ -84,11 +84,11 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { sql(s"set $key=") assert(spark.conf.get(key, "0") === "") - spark.sessionState.conf.clear() + sqlConf.clear() } test("set command for display") { - spark.sessionState.conf.clear() + sqlConf.clear() checkAnswer( sql("SET").where("key = 'spark.sql.groupByOrdinal'").select("key", "value"), Nil) @@ -109,11 +109,11 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("deprecated property") { - spark.sessionState.conf.clear() - val original = spark.conf.get(SQLConf.SHUFFLE_PARTITIONS) + sqlConf.clear() + val original = sqlConf.getConf(SQLConf.SHUFFLE_PARTITIONS) try { sql(s"set ${SQLConf.Deprecated.MAPRED_REDUCE_TASKS}=10") - assert(spark.conf.get(SQLConf.SHUFFLE_PARTITIONS) === 10) + assert(sqlConf.getConf(SQLConf.SHUFFLE_PARTITIONS) === 10) } finally { sql(s"set ${SQLConf.SHUFFLE_PARTITIONS.key}=$original") } @@ -146,18 +146,18 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("reset - public conf") { - spark.sessionState.conf.clear() - val original = spark.conf.get(SQLConf.GROUP_BY_ORDINAL) + sqlConf.clear() + val original = sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL) try { - assert(spark.conf.get(SQLConf.GROUP_BY_ORDINAL)) + assert(sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL)) sql(s"set ${SQLConf.GROUP_BY_ORDINAL.key}=false") - assert(spark.conf.get(SQLConf.GROUP_BY_ORDINAL) === false) + assert(sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL) === false) assert(sql(s"set").where(s"key = '${SQLConf.GROUP_BY_ORDINAL.key}'").count() == 1) - assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES).isEmpty) + assert(sqlConf.getConf(SQLConf.OPTIMIZER_EXCLUDED_RULES).isEmpty) sql(s"reset") - assert(spark.conf.get(SQLConf.GROUP_BY_ORDINAL)) + assert(sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL)) assert(sql(s"set").where(s"key = '${SQLConf.GROUP_BY_ORDINAL.key}'").count() == 0) - assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES) === + assert(sqlConf.getConf(SQLConf.OPTIMIZER_EXCLUDED_RULES) === Some("org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation")) } finally { sql(s"set ${SQLConf.GROUP_BY_ORDINAL.key}=$original") @@ -165,15 +165,15 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("reset - internal conf") { - spark.sessionState.conf.clear() - val original = spark.conf.get(SQLConf.OPTIMIZER_MAX_ITERATIONS) + sqlConf.clear() + val original = sqlConf.getConf(SQLConf.OPTIMIZER_MAX_ITERATIONS) try { - assert(spark.conf.get(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 100) + assert(sqlConf.getConf(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 100) sql(s"set ${SQLConf.OPTIMIZER_MAX_ITERATIONS.key}=10") - assert(spark.conf.get(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 10) + assert(sqlConf.getConf(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 10) assert(sql(s"set").where(s"key = '${SQLConf.OPTIMIZER_MAX_ITERATIONS.key}'").count() == 1) sql(s"reset") - assert(spark.conf.get(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 100) + assert(sqlConf.getConf(SQLConf.OPTIMIZER_MAX_ITERATIONS) === 100) assert(sql(s"set").where(s"key = '${SQLConf.OPTIMIZER_MAX_ITERATIONS.key}'").count() == 0) } finally { sql(s"set ${SQLConf.OPTIMIZER_MAX_ITERATIONS.key}=$original") @@ -181,7 +181,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("reset - user-defined conf") { - spark.sessionState.conf.clear() + sqlConf.clear() val userDefinedConf = "x.y.z.reset" try { assert(spark.conf.getOption(userDefinedConf).isEmpty) @@ -196,7 +196,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("SPARK-32406: reset - single configuration") { - spark.sessionState.conf.clear() + sqlConf.clear() // spark core conf w/o entry registered val appId = spark.sparkContext.getConf.getAppId sql("RESET spark.app.id") @@ -216,19 +216,19 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { sql("RESET spark.abc") // ignore nonexistent keys // runtime sql configs - val original = spark.conf.get(SQLConf.GROUP_BY_ORDINAL) + val original = sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL) sql(s"SET ${SQLConf.GROUP_BY_ORDINAL.key}=false") sql(s"RESET ${SQLConf.GROUP_BY_ORDINAL.key}") - assert(spark.conf.get(SQLConf.GROUP_BY_ORDINAL) === original) + assert(sqlConf.getConf(SQLConf.GROUP_BY_ORDINAL) === original) // runtime sql configs with optional defaults - assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES).isEmpty) + assert(sqlConf.getConf(SQLConf.OPTIMIZER_EXCLUDED_RULES).isEmpty) sql(s"RESET ${SQLConf.OPTIMIZER_EXCLUDED_RULES.key}") - assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES) === + assert(sqlConf.getConf(SQLConf.OPTIMIZER_EXCLUDED_RULES) === Some("org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation")) sql(s"SET ${SQLConf.PLAN_CHANGE_LOG_RULES.key}=abc") sql(s"RESET ${SQLConf.PLAN_CHANGE_LOG_RULES.key}") - assert(spark.conf.get(SQLConf.PLAN_CHANGE_LOG_RULES).isEmpty) + assert(sqlConf.getConf(SQLConf.PLAN_CHANGE_LOG_RULES).isEmpty) // static sql configs checkError( @@ -247,19 +247,19 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("Test ADVISORY_PARTITION_SIZE_IN_BYTES's method") { - spark.sessionState.conf.clear() + sqlConf.clear() spark.conf.set(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key, "100") - assert(spark.conf.get(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 100) + assert(sqlConf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 100) spark.conf.set(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key, "1k") - assert(spark.conf.get(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1024) + assert(sqlConf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1024) spark.conf.set(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key, "1M") - assert(spark.conf.get(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1048576) + assert(sqlConf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1048576) spark.conf.set(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key, "1g") - assert(spark.conf.get(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1073741824) + assert(sqlConf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES) === 1073741824) // test negative value intercept[IllegalArgumentException] { @@ -277,7 +277,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { spark.conf.set(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key, "-90000000000g") } - spark.sessionState.conf.clear() + sqlConf.clear() } test("SparkSession can access configs set in SparkConf") { @@ -305,7 +305,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { try { sparkContext.conf.set(GLOBAL_TEMP_DATABASE, "a") val newSession = new SparkSession(sparkContext) - assert(newSession.conf.get(GLOBAL_TEMP_DATABASE) == "a") + assert(newSession.sessionState.conf.getConf(GLOBAL_TEMP_DATABASE) == "a") checkAnswer( newSession.sql(s"SET ${GLOBAL_TEMP_DATABASE.key}"), Row(GLOBAL_TEMP_DATABASE.key, "a")) @@ -338,16 +338,16 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { } test("SPARK-10365: PARQUET_OUTPUT_TIMESTAMP_TYPE") { - spark.sessionState.conf.clear() + sqlConf.clear() // check default value assert(spark.sessionState.conf.parquetOutputTimestampType == SQLConf.ParquetOutputTimestampType.INT96) - spark.conf.set(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE, "timestamp_micros") + sqlConf.setConf(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE, "timestamp_micros") assert(spark.sessionState.conf.parquetOutputTimestampType == SQLConf.ParquetOutputTimestampType.TIMESTAMP_MICROS) - spark.conf.set(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE, "int96") + sqlConf.setConf(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE, "int96") assert(spark.sessionState.conf.parquetOutputTimestampType == SQLConf.ParquetOutputTimestampType.INT96) @@ -356,7 +356,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { spark.conf.set(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key, "invalid") } - spark.sessionState.conf.clear() + sqlConf.clear() } test("SPARK-22779: correctly compute default value for fallback configs") { @@ -373,10 +373,10 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { .get assert(displayValue === fallback.defaultValueString) - spark.conf.set(SQLConf.PARQUET_COMPRESSION, GZIP.lowerCaseName()) + sqlConf.setConf(SQLConf.PARQUET_COMPRESSION, GZIP.lowerCaseName()) assert(spark.conf.get(fallback.key) === GZIP.lowerCaseName()) - spark.conf.set(fallback, LZO.lowerCaseName()) + sqlConf.setConf(fallback, LZO.lowerCaseName()) assert(spark.conf.get(fallback.key) === LZO.lowerCaseName()) val newDisplayValue = spark.sessionState.conf.getAllDefinedConfs @@ -459,10 +459,10 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { test("set time zone") { TimeZone.getAvailableIDs().foreach { zid => sql(s"set time zone '$zid'") - assert(spark.conf.get(SQLConf.SESSION_LOCAL_TIMEZONE) === zid) + assert(sqlConf.getConf(SQLConf.SESSION_LOCAL_TIMEZONE) === zid) } sql("set time zone local") - assert(spark.conf.get(SQLConf.SESSION_LOCAL_TIMEZONE) === TimeZone.getDefault.getID) + assert(sqlConf.getConf(SQLConf.SESSION_LOCAL_TIMEZONE) === TimeZone.getDefault.getID) val tz = "Invalid TZ" checkError( @@ -476,7 +476,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { (-18 to 18).map(v => (v, s"interval '$v' hours")).foreach { case (i, interval) => sql(s"set time zone $interval") - val zone = spark.conf.get(SQLConf.SESSION_LOCAL_TIMEZONE) + val zone = sqlConf.getConf(SQLConf.SESSION_LOCAL_TIMEZONE) if (i == 0) { assert(zone === "Z") } else { @@ -504,7 +504,7 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { test("SPARK-47765: set collation") { Seq("UNICODE", "UNICODE_CI", "utf8_lcase", "utf8_binary").foreach { collation => sql(s"set collation $collation") - assert(spark.conf.get(SQLConf.DEFAULT_COLLATION) === collation.toUpperCase(Locale.ROOT)) + assert(sqlConf.getConf(SQLConf.DEFAULT_COLLATION) === collation.toUpperCase(Locale.ROOT)) } checkError( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala index 8b11e0c69fa70..24732223c6698 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -54,7 +54,7 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti protected override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(SQLConf.LEGACY_BUCKETED_TABLE_SCAN_OUTPUT_ORDERING, true) + spark.conf.set(SQLConf.LEGACY_BUCKETED_TABLE_SCAN_OUTPUT_ORDERING.key, true) } protected override def afterAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala index 168b6b8629926..e27ec32e287e8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala @@ -50,7 +50,7 @@ abstract class FileStreamSinkSuite extends StreamTest { override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(SQLConf.ORC_IMPLEMENTATION, "native") + spark.conf.set(SQLConf.ORC_IMPLEMENTATION.key, "native") } override def afterAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSourceSuite.scala index 56c4aecb23770..773be0cc08e3f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSourceSuite.scala @@ -262,7 +262,7 @@ class FileStreamSourceSuite extends FileStreamSourceTest { override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set(SQLConf.ORC_IMPLEMENTATION, "native") + spark.conf.set(SQLConf.ORC_IMPLEMENTATION.key, "native") } override def afterAll(): Unit = { @@ -1504,7 +1504,7 @@ class FileStreamSourceSuite extends FileStreamSourceTest { // This is to avoid running a spark job to list of files in parallel // by the InMemoryFileIndex. - spark.conf.set(SQLConf.PARALLEL_PARTITION_DISCOVERY_THRESHOLD, numFiles * 2) + spark.conf.set(SQLConf.PARALLEL_PARTITION_DISCOVERY_THRESHOLD.key, numFiles * 2) withTempDirs { case (root, tmp) => val src = new File(root, "a=1") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala index f3ef73c6af5fa..f7ff39622ed40 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala @@ -1163,7 +1163,7 @@ class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest { func: (Int, Iterator[Int], GroupState[Int]) => Iterator[Int], timeoutType: GroupStateTimeout = GroupStateTimeout.NoTimeout, batchTimestampMs: Long = NO_TIMESTAMP): FlatMapGroupsWithStateExec = { - val stateFormatVersion = spark.conf.get(SQLConf.FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION) + val stateFormatVersion = sqlConf.getConf(SQLConf.FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION) val emptyRdd = spark.sparkContext.emptyRDD[InternalRow] MemoryStream[Int] .toDS() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala index 7ab45e25799bc..68436c4e355b7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala @@ -542,10 +542,7 @@ trait StreamTest extends QueryTest with SharedSparkSession with TimeLimits with val metadataRoot = Option(checkpointLocation).getOrElse(defaultCheckpointLocation) additionalConfs.foreach(pair => { - val value = - if (sparkSession.conf.contains(pair._1)) { - Some(sparkSession.conf.get(pair._1)) - } else None + val value = sparkSession.conf.getOption(pair._1) resetConfValues(pair._1) = value sparkSession.conf.set(pair._1, pair._2) }) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TriggerAvailableNowSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TriggerAvailableNowSuite.scala index defd5fd110de6..a47c2f839692c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TriggerAvailableNowSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TriggerAvailableNowSuite.scala @@ -265,7 +265,7 @@ class TriggerAvailableNowSuite extends FileStreamSourceTest { private def assertQueryUsingRightBatchExecutor( testSource: TestDataFrameProvider, query: StreamingQuery): Unit = { - val useWrapper = query.sparkSession.conf.get( + val useWrapper = query.sparkSession.sessionState.conf.getConf( SQLConf.STREAMING_TRIGGER_AVAILABLE_NOW_WRAPPER_ENABLED) if (useWrapper) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala index ff1473fea369b..4d4cc44eb3e72 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala @@ -103,6 +103,8 @@ trait SharedSparkSessionBase new TestSparkSession(sparkConf) } + protected def sqlConf: SQLConf = _spark.sessionState.conf + /** * Initialize the [[TestSparkSession]]. Generally, this is just called from * beforeAll; however, in test using styles other than FunSuite, there is diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSharedStateSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSharedStateSuite.scala index d84b9f7960231..8c6113fb5569d 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSharedStateSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSharedStateSuite.scala @@ -86,7 +86,7 @@ class HiveSharedStateSuite extends SparkFunSuite { assert(ss2.sparkContext.hadoopConfiguration.get("hive.metastore.warehouse.dir") !== invalidPath, "warehouse conf in session options can't affect application wide hadoop conf") assert(ss.conf.get("spark.foo") === "bar2222", "session level conf should be passed to catalog") - assert(!ss.conf.get(WAREHOUSE_PATH).contains(invalidPath), + assert(!ss.conf.get(WAREHOUSE_PATH.key).contains(invalidPath), "session level conf should be passed to catalog") } diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala index 69abb1d1673ed..865ce81e151c2 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala @@ -828,7 +828,7 @@ object SPARK_18360 { .enableHiveSupport().getOrCreate() val defaultDbLocation = spark.catalog.getDatabase("default").locationUri - assert(new Path(defaultDbLocation) == new Path(spark.conf.get(WAREHOUSE_PATH))) + assert(new Path(defaultDbLocation) == new Path(spark.conf.get(WAREHOUSE_PATH.key))) val hiveClient = spark.sharedState.externalCatalog.unwrapped.asInstanceOf[HiveExternalCatalog].client diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveSerDeReadWriteSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveSerDeReadWriteSuite.scala index aafc4764d2465..1922144a92efa 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveSerDeReadWriteSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveSerDeReadWriteSuite.scala @@ -44,7 +44,7 @@ class HiveSerDeReadWriteSuite extends QueryTest with SQLTestUtils with TestHiveS super.beforeAll() originalConvertMetastoreParquet = spark.conf.get(CONVERT_METASTORE_PARQUET.key) originalConvertMetastoreORC = spark.conf.get(CONVERT_METASTORE_ORC.key) - originalORCImplementation = spark.conf.get(ORC_IMPLEMENTATION) + originalORCImplementation = spark.conf.get(ORC_IMPLEMENTATION.key) spark.conf.set(CONVERT_METASTORE_PARQUET.key, "false") spark.conf.set(CONVERT_METASTORE_ORC.key, "false") diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala index 3deb355e0e4a9..594c097de2c7d 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala @@ -79,13 +79,13 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi test("query global temp view") { val df = Seq(1).toDF("i1") df.createGlobalTempView("tbl1") - val global_temp_db = spark.conf.get(GLOBAL_TEMP_DATABASE) + val global_temp_db = spark.conf.get(GLOBAL_TEMP_DATABASE.key) checkAnswer(spark.sql(s"select * from ${global_temp_db}.tbl1"), Row(1)) spark.sql(s"drop view ${global_temp_db}.tbl1") } test("non-existent global temp view") { - val global_temp_db = spark.conf.get(GLOBAL_TEMP_DATABASE) + val global_temp_db = spark.conf.get(GLOBAL_TEMP_DATABASE.key) val e = intercept[AnalysisException] { spark.sql(s"select * from ${global_temp_db}.nonexistentview") } From a7f191ba5947075066154a33da7908b24c412ccb Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 18 Sep 2024 08:44:22 +0800 Subject: [PATCH 048/189] [SPARK-49640][PS] Apply reservoir sampling in `SampledPlotBase` ### What changes were proposed in this pull request? Apply reservoir sampling in `SampledPlotBase` ### Why are the changes needed? Existing sampling approach has two drawbacks: 1, it needs two jobs to sample `max_rows` rows: - df.count() to compute `fraction = max_rows / count` - df.sample(fraction).to_pandas() to do the sampling 2, the df.sample is based on Bernoulli sampling which **cannot** guarantee the sampled size == expected `max_rows`, e.g. ``` In [1]: df = spark.range(10000) In [2]: [df.sample(0.01).count() for i in range(0, 10)] Out[2]: [96, 97, 95, 97, 105, 105, 105, 87, 95, 110] ``` The size of sampled data is floating near the target size 10000*0.01=100. This relative deviation cannot be ignored, when the input dataset is large and the sampling fraction is small. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? CI and manually check ### Was this patch authored or co-authored using generative AI tooling? No Closes #48105 from zhengruifeng/ps_sampling. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- python/pyspark/pandas/plot/core.py | 51 ++++++++++++++++++++++++------ 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index 067c7db664dee..7630ecc398954 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -68,19 +68,52 @@ class SampledPlotBase: def get_sampled(self, data): from pyspark.pandas import DataFrame, Series + if not isinstance(data, (DataFrame, Series)): + raise TypeError("Only DataFrame and Series are supported for plotting.") + if isinstance(data, Series): + data = data.to_frame() + fraction = get_option("plotting.sample_ratio") - if fraction is None: - fraction = 1 / (len(data) / get_option("plotting.max_rows")) - fraction = min(1.0, fraction) - self.fraction = fraction - - if isinstance(data, (DataFrame, Series)): - if isinstance(data, Series): - data = data.to_frame() + if fraction is not None: + self.fraction = fraction sampled = data._internal.resolved_copy.spark_frame.sample(fraction=self.fraction) return DataFrame(data._internal.with_new_sdf(sampled))._to_pandas() else: - raise TypeError("Only DataFrame and Series are supported for plotting.") + from pyspark.sql import Observation + + max_rows = get_option("plotting.max_rows") + observation = Observation("ps plotting") + sdf = data._internal.resolved_copy.spark_frame.observe( + observation, F.count(F.lit(1)).alias("count") + ) + + rand_col_name = "__ps_plotting_sampled_plot_base_rand__" + id_col_name = "__ps_plotting_sampled_plot_base_id__" + + sampled = ( + sdf.select( + "*", + F.rand().alias(rand_col_name), + F.monotonically_increasing_id().alias(id_col_name), + ) + .sort(rand_col_name) + .limit(max_rows + 1) + .coalesce(1) + .sortWithinPartitions(id_col_name) + .drop(rand_col_name, id_col_name) + ) + + pdf = DataFrame(data._internal.with_new_sdf(sampled))._to_pandas() + + if len(pdf) > max_rows: + try: + self.fraction = float(max_rows) / observation.get["count"] + except Exception: + pass + return pdf[:max_rows] + else: + self.fraction = 1.0 + return pdf def set_result_text(self, ax): assert hasattr(self, "fraction") From b1807095bef9c6d98e60bdc2669c8af93bc68ad4 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Wed, 18 Sep 2024 10:26:35 +0800 Subject: [PATCH 049/189] [SPARK-49495][DOCS] Document and Feature Preview on the master branch via Live GitHub Pages Updates ### What changes were proposed in this pull request? This pull request introduces functionalities that enable 'Document and Feature Preview on the master branch via Live GitHub Pages Updates'. ### Why are the changes needed? - Instead of limited 72-hour voting phases, it provides the developer community with more opportunities to preview and verify the documentation contents. - Instead of waiting for the final announcement of an official spark feature release, users can now preview some of the ongoing documented features, increasing the willingness to upgrade, sensing breaking changes in advance, and reducing the burden during the final upgrades. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? An [example](https://yaooqinn.github.io/spark-gh-pages/) has been established by this repo - https://github.com/yaooqinn/spark-gh-pages - Broken logo links are fixed at https://github.com/apache/spark/pull/47966 ### Was this patch authored or co-authored using generative AI tooling? no Closes #47968 from yaooqinn/SPARK-49495. Authored-by: Kent Yao Signed-off-by: Kent Yao --- .github/workflows/pages.yml | 90 +++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 .github/workflows/pages.yml diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml new file mode 100644 index 0000000000000..083620427c015 --- /dev/null +++ b/.github/workflows/pages.yml @@ -0,0 +1,90 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +name: GitHub Pages deployment + +on: + push: + branches: + - master + +concurrency: + group: 'docs preview' + cancel-in-progress: true + +jobs: + docs: + name: Build and deploy documentation + runs-on: ubuntu-latest + permissions: + id-token: write + pages: write + env: + SPARK_TESTING: 1 # Reduce some noise in the logs + RELEASE_VERSION: 'In-Progress' + steps: + - name: Checkout Spark repository + uses: actions/checkout@v4 + with: + repository: apache/spark + ref: 'master' + - name: Install Java 17 + uses: actions/setup-java@v4 + with: + distribution: zulu + java-version: 17 + - name: Install Python 3.9 + uses: actions/setup-python@v5 + with: + python-version: '3.9' + architecture: x64 + cache: 'pip' + - name: Install Python dependencies + run: pip install --upgrade -r dev/requirements.txt + - name: Install Ruby for documentation generation + uses: ruby/setup-ruby@v1 + with: + ruby-version: '3.3' + bundler-cache: true + - name: Install Pandoc + uses: pandoc/actions/setup@d6abb76f6c8a1a9a5e15a5190c96a02aabffd1ee + with: + version: 3.3 + - name: Install dependencies for documentation generation + run: | + cd docs + gem install bundler -v 2.4.22 -n /usr/local/bin + bundle install --retry=100 + - name: Run documentation build + run: | + sed -i".tmp1" 's/SPARK_VERSION:.*$/SPARK_VERSION: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp2" 's/SPARK_VERSION_SHORT:.*$/SPARK_VERSION_SHORT: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp3" "s/'facetFilters':.*$/'facetFilters': [\"version:$RELEASE_VERSION\"]/g" docs/_config.yml + sed -i".tmp4" 's/__version__: str = .*$/__version__: str = "'"$RELEASE_VERSION"'"/' python/pyspark/version.py + cd docs + SKIP_RDOC=1 bundle exec jekyll build + - name: Setup Pages + uses: actions/configure-pages@v5 + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: 'docs/_site' + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 From 31cd6991558c45cea56ba25cb89f13e64e3d93fa Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Tue, 17 Sep 2024 23:00:49 -0400 Subject: [PATCH 050/189] [SPARK-49424][CONNECT][SQL] Consolidate Encoders.scala ### What changes were proposed in this pull request? This PR moves Encoders.scala to sql/api. It removes the duplicate one in connect. ### Why are the changes needed? We are creating a unified scala interface for Classic and Connect. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48021 from hvanhovell/SPARK-49424. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../resources/error/error-conditions.json | 6 + project/MimaExcludes.scala | 4 + .../scala/org/apache/spark/sql/Encoders.scala | 152 +++++--- .../catalyst/encoders/AgnosticEncoder.scala | 13 + .../spark/sql/errors/ExecutionErrors.scala | 18 + .../scala/org/apache/spark/sql/Encoders.scala | 348 ------------------ .../catalyst/encoders/ExpressionEncoder.scala | 51 +-- .../sql/errors/QueryExecutionErrors.scala | 11 - .../encoders/EncoderResolutionSuite.scala | 8 +- .../encoders/ExpressionEncoderSuite.scala | 23 +- .../connect/planner/SparkConnectPlanner.scala | 17 +- .../scala/org/apache/spark/sql/Dataset.scala | 84 +++-- .../spark/sql/KeyValueGroupedDataset.scala | 37 +- .../spark/sql/RelationalGroupedDataset.scala | 8 +- .../spark/sql/execution/aggregate/udaf.scala | 2 +- .../ContinuousTextSocketSource.scala | 6 +- .../sql/expressions/ReduceAggregator.scala | 11 +- .../spark/sql/internal/TypedAggUtils.scala | 8 +- ...latMapGroupsWithStateExecHelperSuite.scala | 4 +- .../streaming/state/ListStateSuite.scala | 18 +- .../streaming/state/MapStateSuite.scala | 15 +- .../state/StatefulProcessorHandleSuite.scala | 24 +- .../streaming/state/TimerSuite.scala | 22 +- .../streaming/state/ValueStateSuite.scala | 30 +- 24 files changed, 300 insertions(+), 620 deletions(-) rename {connector/connect/client/jvm => sql/api}/src/main/scala/org/apache/spark/sql/Encoders.scala (77%) delete mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/Encoders.scala diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 57b3d33741e98..25dd676c4aff9 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1996,6 +1996,12 @@ }, "sqlState" : "42903" }, + "INVALID_AGNOSTIC_ENCODER" : { + "message" : [ + "Found an invalid agnostic encoder. Expects an instance of AgnosticEncoder but got . For more information consult '/api/java/index.html?org/apache/spark/sql/Encoder.html'." + ], + "sqlState" : "42001" + }, "INVALID_ARRAY_INDEX" : { "message" : [ "The index is out of bounds. The array has elements. Use the SQL function `get()` to tolerate accessing element at invalid index and return NULL instead. If necessary set to \"false\" to bypass this error." diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 68433b501bcc4..dfe7b14e2ec66 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -160,6 +160,10 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.DataFrameWriterV2"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.WriteConfigMethods"), + // SPARK-49424: Shared Encoders + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.Encoders"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.Encoders$"), + // SPARK-49413: Create a shared RuntimeConfig interface. ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.RuntimeConfig"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.RuntimeConfig$"), diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Encoders.scala b/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala similarity index 77% rename from connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Encoders.scala rename to sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala index 33a322109c1b6..9976b34f7a01f 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Encoders.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala @@ -14,95 +14,99 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + package org.apache.spark.sql -import scala.reflect.ClassTag +import java.lang.reflect.Modifier + +import scala.reflect.{classTag, ClassTag} import scala.reflect.runtime.universe.TypeTag import org.apache.spark.sql.catalyst.{JavaTypeInference, ScalaReflection} -import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, JavaSerializationCodec, KryoSerializationCodec, RowEncoder => RowEncoderFactory} +import org.apache.spark.sql.catalyst.encoders.{Codec, JavaSerializationCodec, KryoSerializationCodec, RowEncoder => SchemaInference} import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders._ -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.errors.ExecutionErrors +import org.apache.spark.sql.types._ /** * Methods for creating an [[Encoder]]. * - * @since 3.5.0 + * @since 1.6.0 */ object Encoders { /** * An encoder for nullable boolean type. The Scala primitive encoder is available as * [[scalaBoolean]]. - * @since 3.5.0 + * @since 1.6.0 */ def BOOLEAN: Encoder[java.lang.Boolean] = BoxedBooleanEncoder /** * An encoder for nullable byte type. The Scala primitive encoder is available as [[scalaByte]]. - * @since 3.5.0 + * @since 1.6.0 */ def BYTE: Encoder[java.lang.Byte] = BoxedByteEncoder /** * An encoder for nullable short type. The Scala primitive encoder is available as * [[scalaShort]]. - * @since 3.5.0 + * @since 1.6.0 */ def SHORT: Encoder[java.lang.Short] = BoxedShortEncoder /** * An encoder for nullable int type. The Scala primitive encoder is available as [[scalaInt]]. - * @since 3.5.0 + * @since 1.6.0 */ def INT: Encoder[java.lang.Integer] = BoxedIntEncoder /** * An encoder for nullable long type. The Scala primitive encoder is available as [[scalaLong]]. - * @since 3.5.0 + * @since 1.6.0 */ def LONG: Encoder[java.lang.Long] = BoxedLongEncoder /** * An encoder for nullable float type. The Scala primitive encoder is available as * [[scalaFloat]]. - * @since 3.5.0 + * @since 1.6.0 */ def FLOAT: Encoder[java.lang.Float] = BoxedFloatEncoder /** * An encoder for nullable double type. The Scala primitive encoder is available as * [[scalaDouble]]. - * @since 3.5.0 + * @since 1.6.0 */ def DOUBLE: Encoder[java.lang.Double] = BoxedDoubleEncoder /** * An encoder for nullable string type. * - * @since 3.5.0 + * @since 1.6.0 */ def STRING: Encoder[java.lang.String] = StringEncoder /** * An encoder for nullable decimal type. * - * @since 3.5.0 + * @since 1.6.0 */ def DECIMAL: Encoder[java.math.BigDecimal] = DEFAULT_JAVA_DECIMAL_ENCODER /** * An encoder for nullable date type. * - * @since 3.5.0 + * @since 1.6.0 */ - def DATE: Encoder[java.sql.Date] = DateEncoder(lenientSerialization = false) + def DATE: Encoder[java.sql.Date] = STRICT_DATE_ENCODER /** * Creates an encoder that serializes instances of the `java.time.LocalDate` class to the * internal representation of nullable Catalyst's DateType. * - * @since 3.5.0 + * @since 3.0.0 */ def LOCALDATE: Encoder[java.time.LocalDate] = STRICT_LOCAL_DATE_ENCODER @@ -110,14 +114,14 @@ object Encoders { * Creates an encoder that serializes instances of the `java.time.LocalDateTime` class to the * internal representation of nullable Catalyst's TimestampNTZType. * - * @since 3.5.0 + * @since 3.4.0 */ def LOCALDATETIME: Encoder[java.time.LocalDateTime] = LocalDateTimeEncoder /** * An encoder for nullable timestamp type. * - * @since 3.5.0 + * @since 1.6.0 */ def TIMESTAMP: Encoder[java.sql.Timestamp] = STRICT_TIMESTAMP_ENCODER @@ -125,14 +129,14 @@ object Encoders { * Creates an encoder that serializes instances of the `java.time.Instant` class to the internal * representation of nullable Catalyst's TimestampType. * - * @since 3.5.0 + * @since 3.0.0 */ def INSTANT: Encoder[java.time.Instant] = STRICT_INSTANT_ENCODER /** * An encoder for arrays of bytes. * - * @since 3.5.0 + * @since 1.6.1 */ def BINARY: Encoder[Array[Byte]] = BinaryEncoder @@ -140,7 +144,7 @@ object Encoders { * Creates an encoder that serializes instances of the `java.time.Duration` class to the * internal representation of nullable Catalyst's DayTimeIntervalType. * - * @since 3.5.0 + * @since 3.2.0 */ def DURATION: Encoder[java.time.Duration] = DayTimeIntervalEncoder @@ -148,7 +152,7 @@ object Encoders { * Creates an encoder that serializes instances of the `java.time.Period` class to the internal * representation of nullable Catalyst's YearMonthIntervalType. * - * @since 3.5.0 + * @since 3.2.0 */ def PERIOD: Encoder[java.time.Period] = YearMonthIntervalEncoder @@ -166,7 +170,7 @@ object Encoders { * - collection types: array, java.util.List, and map * - nested java bean. * - * @since 3.5.0 + * @since 1.6.0 */ def bean[T](beanClass: Class[T]): Encoder[T] = JavaTypeInference.encoderFor(beanClass) @@ -175,71 +179,96 @@ object Encoders { * * @since 3.5.0 */ - def row(schema: StructType): Encoder[Row] = RowEncoderFactory.encoderFor(schema) + def row(schema: StructType): Encoder[Row] = SchemaInference.encoderFor(schema) /** - * (Scala-specific) Creates an encoder that serializes objects of type T using generic Java - * serialization. This encoder maps T into a single byte array (binary) field. + * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo. This + * encoder maps T into a single byte array (binary) field. * * T must be publicly accessible. * - * @note - * This is extremely inefficient and should only be used as the last resort. - * @since 4.0.0 + * @since 1.6.0 */ - def javaSerialization[T: ClassTag]: Encoder[T] = { - TransformingEncoder(implicitly[ClassTag[T]], BinaryEncoder, JavaSerializationCodec) - } + def kryo[T: ClassTag]: Encoder[T] = genericSerializer(KryoSerializationCodec) /** - * Creates an encoder that serializes objects of type T using generic Java serialization. This - * encoder maps T into a single byte array (binary) field. + * Creates an encoder that serializes objects of type T using Kryo. This encoder maps T into a + * single byte array (binary) field. + * + * T must be publicly accessible. + * + * @since 1.6.0 + */ + def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz)) + + /** + * (Scala-specific) Creates an encoder that serializes objects of type T using generic Java + * serialization. This encoder maps T into a single byte array (binary) field. * * T must be publicly accessible. * * @note * This is extremely inefficient and should only be used as the last resort. - * @since 4.0.0 + * + * @since 1.6.0 */ - def javaSerialization[T](clazz: Class[T]): Encoder[T] = javaSerialization(ClassTag[T](clazz)) + def javaSerialization[T: ClassTag]: Encoder[T] = genericSerializer(JavaSerializationCodec) /** - * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo. This + * Creates an encoder that serializes objects of type T using generic Java serialization. This * encoder maps T into a single byte array (binary) field. * * T must be publicly accessible. * - * @since 4.0.0 + * @note + * This is extremely inefficient and should only be used as the last resort. + * + * @since 1.6.0 */ - def kryo[T: ClassTag]: Encoder[T] = { - TransformingEncoder(implicitly[ClassTag[T]], BinaryEncoder, KryoSerializationCodec) + def javaSerialization[T](clazz: Class[T]): Encoder[T] = + javaSerialization(ClassTag[T](clazz)) + + /** Throws an exception if T is not a public class. */ + private def validatePublicClass[T: ClassTag](): Unit = { + if (!Modifier.isPublic(classTag[T].runtimeClass.getModifiers)) { + throw ExecutionErrors.notPublicClassError(classTag[T].runtimeClass.getName) + } + } + + /** A way to construct encoders using generic serializers. */ + private def genericSerializer[T: ClassTag]( + provider: () => Codec[Any, Array[Byte]]): Encoder[T] = { + if (classTag[T].runtimeClass.isPrimitive) { + throw ExecutionErrors.primitiveTypesNotSupportedError() + } + + validatePublicClass[T]() + + TransformingEncoder(classTag[T], BinaryEncoder, provider) + } + + private[sql] def tupleEncoder[T](encoders: Encoder[_]*): Encoder[T] = { + ProductEncoder.tuple(encoders.map(agnosticEncoderFor(_))).asInstanceOf[Encoder[T]] } /** - * Creates an encoder that serializes objects of type T using Kryo. This encoder maps T into a - * single byte array (binary) field. - * - * T must be publicly accessible. + * An encoder for 1-ary tuples. * * @since 4.0.0 */ - def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz)) - - private def tupleEncoder[T](encoders: Encoder[_]*): Encoder[T] = { - ProductEncoder.tuple(encoders.asInstanceOf[Seq[AgnosticEncoder[_]]]).asInstanceOf[Encoder[T]] - } + def tuple[T1](e1: Encoder[T1]): Encoder[(T1)] = tupleEncoder(e1) /** * An encoder for 2-ary tuples. * - * @since 3.5.0 + * @since 1.6.0 */ def tuple[T1, T2](e1: Encoder[T1], e2: Encoder[T2]): Encoder[(T1, T2)] = tupleEncoder(e1, e2) /** * An encoder for 3-ary tuples. * - * @since 3.5.0 + * @since 1.6.0 */ def tuple[T1, T2, T3]( e1: Encoder[T1], @@ -249,7 +278,7 @@ object Encoders { /** * An encoder for 4-ary tuples. * - * @since 3.5.0 + * @since 1.6.0 */ def tuple[T1, T2, T3, T4]( e1: Encoder[T1], @@ -260,7 +289,7 @@ object Encoders { /** * An encoder for 5-ary tuples. * - * @since 3.5.0 + * @since 1.6.0 */ def tuple[T1, T2, T3, T4, T5]( e1: Encoder[T1], @@ -271,49 +300,50 @@ object Encoders { /** * An encoder for Scala's product type (tuples, case classes, etc). - * @since 3.5.0 + * @since 2.0.0 */ def product[T <: Product: TypeTag]: Encoder[T] = ScalaReflection.encoderFor[T] /** * An encoder for Scala's primitive int type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaInt: Encoder[Int] = PrimitiveIntEncoder /** * An encoder for Scala's primitive long type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaLong: Encoder[Long] = PrimitiveLongEncoder /** * An encoder for Scala's primitive double type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaDouble: Encoder[Double] = PrimitiveDoubleEncoder /** * An encoder for Scala's primitive float type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaFloat: Encoder[Float] = PrimitiveFloatEncoder /** * An encoder for Scala's primitive byte type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaByte: Encoder[Byte] = PrimitiveByteEncoder /** * An encoder for Scala's primitive short type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaShort: Encoder[Short] = PrimitiveShortEncoder /** * An encoder for Scala's primitive boolean type. - * @since 3.5.0 + * @since 2.0.0 */ def scalaBoolean: Encoder[Boolean] = PrimitiveBooleanEncoder + } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala index a578495755492..10f734b3f84ed 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala @@ -46,7 +46,20 @@ trait AgnosticEncoder[T] extends Encoder[T] { def isStruct: Boolean = false } +/** + * Extract an [[AgnosticEncoder]] from an [[Encoder]]. + */ +trait ToAgnosticEncoder[T] { + def encoder: AgnosticEncoder[T] +} + object AgnosticEncoders { + def agnosticEncoderFor[T: Encoder]: AgnosticEncoder[T] = implicitly[Encoder[T]] match { + case a: AgnosticEncoder[T] => a + case e: ToAgnosticEncoder[T @unchecked] => e.encoder + case other => throw ExecutionErrors.invalidAgnosticEncoderError(other) + } + case class OptionEncoder[E](elementEncoder: AgnosticEncoder[E]) extends AgnosticEncoder[Option[E]] { override def isPrimitive: Boolean = false diff --git a/sql/api/src/main/scala/org/apache/spark/sql/errors/ExecutionErrors.scala b/sql/api/src/main/scala/org/apache/spark/sql/errors/ExecutionErrors.scala index 4890ff4431fe6..698a7b096e1a5 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/errors/ExecutionErrors.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/errors/ExecutionErrors.scala @@ -217,9 +217,27 @@ private[sql] trait ExecutionErrors extends DataTypeErrorsBase { new SparkRuntimeException(errorClass = "CANNOT_USE_KRYO", messageParameters = Map.empty) } + def notPublicClassError(name: String): SparkUnsupportedOperationException = { + new SparkUnsupportedOperationException( + errorClass = "_LEGACY_ERROR_TEMP_2229", + messageParameters = Map("name" -> name)) + } + + def primitiveTypesNotSupportedError(): SparkUnsupportedOperationException = { + new SparkUnsupportedOperationException(errorClass = "_LEGACY_ERROR_TEMP_2230") + } + def elementsOfTupleExceedLimitError(): SparkUnsupportedOperationException = { new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_2150") } + + def invalidAgnosticEncoderError(encoder: AnyRef): Throwable = { + new SparkRuntimeException( + errorClass = "INVALID_AGNOSTIC_ENCODER", + messageParameters = Map( + "encoderType" -> encoder.getClass.getName, + "docroot" -> SparkBuildInfo.spark_doc_root)) + } } private[sql] object ExecutionErrors extends ExecutionErrors diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoders.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoders.scala deleted file mode 100644 index 7e040f6232fbe..0000000000000 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoders.scala +++ /dev/null @@ -1,348 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.spark.sql - -import java.lang.reflect.Modifier - -import scala.reflect.{classTag, ClassTag} -import scala.reflect.runtime.universe.TypeTag - -import org.apache.spark.sql.catalyst.encoders.{encoderFor, Codec, ExpressionEncoder, JavaSerializationCodec, KryoSerializationCodec} -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{BinaryEncoder, TransformingEncoder} -import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.types._ - -/** - * Methods for creating an [[Encoder]]. - * - * @since 1.6.0 - */ -object Encoders { - - /** - * An encoder for nullable boolean type. - * The Scala primitive encoder is available as [[scalaBoolean]]. - * @since 1.6.0 - */ - def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder() - - /** - * An encoder for nullable byte type. - * The Scala primitive encoder is available as [[scalaByte]]. - * @since 1.6.0 - */ - def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder() - - /** - * An encoder for nullable short type. - * The Scala primitive encoder is available as [[scalaShort]]. - * @since 1.6.0 - */ - def SHORT: Encoder[java.lang.Short] = ExpressionEncoder() - - /** - * An encoder for nullable int type. - * The Scala primitive encoder is available as [[scalaInt]]. - * @since 1.6.0 - */ - def INT: Encoder[java.lang.Integer] = ExpressionEncoder() - - /** - * An encoder for nullable long type. - * The Scala primitive encoder is available as [[scalaLong]]. - * @since 1.6.0 - */ - def LONG: Encoder[java.lang.Long] = ExpressionEncoder() - - /** - * An encoder for nullable float type. - * The Scala primitive encoder is available as [[scalaFloat]]. - * @since 1.6.0 - */ - def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder() - - /** - * An encoder for nullable double type. - * The Scala primitive encoder is available as [[scalaDouble]]. - * @since 1.6.0 - */ - def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder() - - /** - * An encoder for nullable string type. - * - * @since 1.6.0 - */ - def STRING: Encoder[java.lang.String] = ExpressionEncoder() - - /** - * An encoder for nullable decimal type. - * - * @since 1.6.0 - */ - def DECIMAL: Encoder[java.math.BigDecimal] = ExpressionEncoder() - - /** - * An encoder for nullable date type. - * - * @since 1.6.0 - */ - def DATE: Encoder[java.sql.Date] = ExpressionEncoder() - - /** - * Creates an encoder that serializes instances of the `java.time.LocalDate` class - * to the internal representation of nullable Catalyst's DateType. - * - * @since 3.0.0 - */ - def LOCALDATE: Encoder[java.time.LocalDate] = ExpressionEncoder() - - /** - * Creates an encoder that serializes instances of the `java.time.LocalDateTime` class - * to the internal representation of nullable Catalyst's TimestampNTZType. - * - * @since 3.4.0 - */ - def LOCALDATETIME: Encoder[java.time.LocalDateTime] = ExpressionEncoder() - - /** - * An encoder for nullable timestamp type. - * - * @since 1.6.0 - */ - def TIMESTAMP: Encoder[java.sql.Timestamp] = ExpressionEncoder() - - /** - * Creates an encoder that serializes instances of the `java.time.Instant` class - * to the internal representation of nullable Catalyst's TimestampType. - * - * @since 3.0.0 - */ - def INSTANT: Encoder[java.time.Instant] = ExpressionEncoder() - - /** - * An encoder for arrays of bytes. - * - * @since 1.6.1 - */ - def BINARY: Encoder[Array[Byte]] = ExpressionEncoder() - - /** - * Creates an encoder that serializes instances of the `java.time.Duration` class - * to the internal representation of nullable Catalyst's DayTimeIntervalType. - * - * @since 3.2.0 - */ - def DURATION: Encoder[java.time.Duration] = ExpressionEncoder() - - /** - * Creates an encoder that serializes instances of the `java.time.Period` class - * to the internal representation of nullable Catalyst's YearMonthIntervalType. - * - * @since 3.2.0 - */ - def PERIOD: Encoder[java.time.Period] = ExpressionEncoder() - - /** - * Creates an encoder for Java Bean of type T. - * - * T must be publicly accessible. - * - * supported types for java bean field: - * - primitive types: boolean, int, double, etc. - * - boxed types: Boolean, Integer, Double, etc. - * - String - * - java.math.BigDecimal, java.math.BigInteger - * - time related: java.sql.Date, java.sql.Timestamp, java.time.LocalDate, java.time.Instant - * - collection types: array, java.util.List, and map - * - nested java bean. - * - * @since 1.6.0 - */ - def bean[T](beanClass: Class[T]): Encoder[T] = ExpressionEncoder.javaBean(beanClass) - - /** - * Creates a [[Row]] encoder for schema `schema`. - * - * @since 3.5.0 - */ - def row(schema: StructType): Encoder[Row] = ExpressionEncoder(schema) - - /** - * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo. - * This encoder maps T into a single byte array (binary) field. - * - * T must be publicly accessible. - * - * @since 1.6.0 - */ - def kryo[T: ClassTag]: Encoder[T] = genericSerializer(KryoSerializationCodec) - - /** - * Creates an encoder that serializes objects of type T using Kryo. - * This encoder maps T into a single byte array (binary) field. - * - * T must be publicly accessible. - * - * @since 1.6.0 - */ - def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz)) - - /** - * (Scala-specific) Creates an encoder that serializes objects of type T using generic Java - * serialization. This encoder maps T into a single byte array (binary) field. - * - * T must be publicly accessible. - * - * @note This is extremely inefficient and should only be used as the last resort. - * - * @since 1.6.0 - */ - def javaSerialization[T: ClassTag]: Encoder[T] = genericSerializer(JavaSerializationCodec) - - /** - * Creates an encoder that serializes objects of type T using generic Java serialization. - * This encoder maps T into a single byte array (binary) field. - * - * T must be publicly accessible. - * - * @note This is extremely inefficient and should only be used as the last resort. - * - * @since 1.6.0 - */ - def javaSerialization[T](clazz: Class[T]): Encoder[T] = javaSerialization(ClassTag[T](clazz)) - - /** Throws an exception if T is not a public class. */ - private def validatePublicClass[T: ClassTag](): Unit = { - if (!Modifier.isPublic(classTag[T].runtimeClass.getModifiers)) { - throw QueryExecutionErrors.notPublicClassError(classTag[T].runtimeClass.getName) - } - } - - /** A way to construct encoders using generic serializers. */ - private def genericSerializer[T: ClassTag]( - provider: () => Codec[Any, Array[Byte]]): Encoder[T] = { - if (classTag[T].runtimeClass.isPrimitive) { - throw QueryExecutionErrors.primitiveTypesNotSupportedError() - } - - validatePublicClass[T]() - - ExpressionEncoder(TransformingEncoder(classTag[T], BinaryEncoder, provider)) - } - - /** - * An encoder for 2-ary tuples. - * - * @since 1.6.0 - */ - def tuple[T1, T2]( - e1: Encoder[T1], - e2: Encoder[T2]): Encoder[(T1, T2)] = { - ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2)) - } - - /** - * An encoder for 3-ary tuples. - * - * @since 1.6.0 - */ - def tuple[T1, T2, T3]( - e1: Encoder[T1], - e2: Encoder[T2], - e3: Encoder[T3]): Encoder[(T1, T2, T3)] = { - ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3)) - } - - /** - * An encoder for 4-ary tuples. - * - * @since 1.6.0 - */ - def tuple[T1, T2, T3, T4]( - e1: Encoder[T1], - e2: Encoder[T2], - e3: Encoder[T3], - e4: Encoder[T4]): Encoder[(T1, T2, T3, T4)] = { - ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3), encoderFor(e4)) - } - - /** - * An encoder for 5-ary tuples. - * - * @since 1.6.0 - */ - def tuple[T1, T2, T3, T4, T5]( - e1: Encoder[T1], - e2: Encoder[T2], - e3: Encoder[T3], - e4: Encoder[T4], - e5: Encoder[T5]): Encoder[(T1, T2, T3, T4, T5)] = { - ExpressionEncoder.tuple( - encoderFor(e1), encoderFor(e2), encoderFor(e3), encoderFor(e4), encoderFor(e5)) - } - - /** - * An encoder for Scala's product type (tuples, case classes, etc). - * @since 2.0.0 - */ - def product[T <: Product : TypeTag]: Encoder[T] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive int type. - * @since 2.0.0 - */ - def scalaInt: Encoder[Int] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive long type. - * @since 2.0.0 - */ - def scalaLong: Encoder[Long] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive double type. - * @since 2.0.0 - */ - def scalaDouble: Encoder[Double] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive float type. - * @since 2.0.0 - */ - def scalaFloat: Encoder[Float] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive byte type. - * @since 2.0.0 - */ - def scalaByte: Encoder[Byte] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive short type. - * @since 2.0.0 - */ - def scalaShort: Encoder[Short] = ExpressionEncoder() - - /** - * An encoder for Scala's primitive boolean type. - * @since 2.0.0 - */ - def scalaBoolean: Encoder[Boolean] = ExpressionEncoder() - -} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala index 8e39ae0389c2c..d7d53230470d9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala @@ -70,54 +70,6 @@ object ExpressionEncoder { apply(JavaTypeInference.encoderFor(beanClass)) } - /** - * Given a set of N encoders, constructs a new encoder that produce objects as items in an - * N-tuple. Note that these encoders should be unresolved so that information about - * name/positional binding is preserved. - * When `useNullSafeDeserializer` is true, the deserialization result for a child will be null if - * the input is null. It is false by default as most deserializers handle null input properly and - * don't require an extra null check. Some of them are null-tolerant, such as the deserializer for - * `Option[T]`, and we must not set it to true in this case. - */ - def tuple( - encoders: Seq[ExpressionEncoder[_]], - useNullSafeDeserializer: Boolean = false): ExpressionEncoder[_] = { - val tupleEncoder = AgnosticEncoders.ProductEncoder.tuple( - encoders.map(_.encoder), - useNullSafeDeserializer) - ExpressionEncoder(tupleEncoder) - } - - // Tuple1 - def tuple[T](e: ExpressionEncoder[T]): ExpressionEncoder[Tuple1[T]] = - tuple(Seq(e)).asInstanceOf[ExpressionEncoder[Tuple1[T]]] - - def tuple[T1, T2]( - e1: ExpressionEncoder[T1], - e2: ExpressionEncoder[T2]): ExpressionEncoder[(T1, T2)] = - tuple(Seq(e1, e2)).asInstanceOf[ExpressionEncoder[(T1, T2)]] - - def tuple[T1, T2, T3]( - e1: ExpressionEncoder[T1], - e2: ExpressionEncoder[T2], - e3: ExpressionEncoder[T3]): ExpressionEncoder[(T1, T2, T3)] = - tuple(Seq(e1, e2, e3)).asInstanceOf[ExpressionEncoder[(T1, T2, T3)]] - - def tuple[T1, T2, T3, T4]( - e1: ExpressionEncoder[T1], - e2: ExpressionEncoder[T2], - e3: ExpressionEncoder[T3], - e4: ExpressionEncoder[T4]): ExpressionEncoder[(T1, T2, T3, T4)] = - tuple(Seq(e1, e2, e3, e4)).asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4)]] - - def tuple[T1, T2, T3, T4, T5]( - e1: ExpressionEncoder[T1], - e2: ExpressionEncoder[T2], - e3: ExpressionEncoder[T3], - e4: ExpressionEncoder[T4], - e5: ExpressionEncoder[T5]): ExpressionEncoder[(T1, T2, T3, T4, T5)] = - tuple(Seq(e1, e2, e3, e4, e5)).asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4, T5)]] - private val anyObjectType = ObjectType(classOf[Any]) /** @@ -189,7 +141,8 @@ case class ExpressionEncoder[T]( encoder: AgnosticEncoder[T], objSerializer: Expression, objDeserializer: Expression) - extends Encoder[T] { + extends Encoder[T] + with ToAgnosticEncoder[T] { override def clsTag: ClassTag[T] = encoder.clsTag diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 2ab86a5c5f03f..4bc071155012b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -1876,17 +1876,6 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE cause = null) } - def notPublicClassError(name: String): SparkUnsupportedOperationException = { - new SparkUnsupportedOperationException( - errorClass = "_LEGACY_ERROR_TEMP_2229", - messageParameters = Map( - "name" -> name)) - } - - def primitiveTypesNotSupportedError(): SparkUnsupportedOperationException = { - new SparkUnsupportedOperationException(errorClass = "_LEGACY_ERROR_TEMP_2230") - } - def onlySupportDataSourcesProvidingFileFormatError(providingClass: String): Throwable = { new SparkException( errorClass = "_LEGACY_ERROR_TEMP_2233", diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderResolutionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderResolutionSuite.scala index 28796db7c02e0..35a27f41da80a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderResolutionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderResolutionSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.catalyst.encoders import scala.reflect.runtime.universe.TypeTag import org.apache.spark.SparkRuntimeException -import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.{AnalysisException, Encoders} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} @@ -71,9 +71,9 @@ class EncoderResolutionSuite extends PlanTest { } test("real type doesn't match encoder schema but they are compatible: tupled encoder") { - val encoder = ExpressionEncoder.tuple( - ExpressionEncoder[StringLongClass](), - ExpressionEncoder[Long]()) + val encoder = encoderFor(Encoders.tuple( + Encoders.product[StringLongClass], + Encoders.scalaLong)) val attrs = Seq($"a".struct($"a".string, $"b".byte), $"b".int) testFromRow(encoder, attrs, InternalRow(InternalRow(str, 1.toByte), 2)) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala index 0c0c7f12f1764..3b5cbed2cc527 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala @@ -321,29 +321,29 @@ class ExpressionEncoderSuite extends CodegenInterpretedPlanTest with AnalysisTes encodeDecodeTest( 1 -> 10L, "tuple with 2 flat encoders")( - ExpressionEncoder.tuple(ExpressionEncoder[Int](), ExpressionEncoder[Long]())) + encoderFor(Encoders.tuple(Encoders.scalaInt, Encoders.scalaLong))) encodeDecodeTest( (PrimitiveData(1, 1, 1, 1, 1, 1, true), (3, 30L)), "tuple with 2 product encoders")( - ExpressionEncoder.tuple(ExpressionEncoder[PrimitiveData](), ExpressionEncoder[(Int, Long)]())) + encoderFor(Encoders.tuple(Encoders.product[PrimitiveData], Encoders.product[(Int, Long)]))) encodeDecodeTest( (PrimitiveData(1, 1, 1, 1, 1, 1, true), 3), "tuple with flat encoder and product encoder")( - ExpressionEncoder.tuple(ExpressionEncoder[PrimitiveData](), ExpressionEncoder[Int]())) + encoderFor(Encoders.tuple(Encoders.product[PrimitiveData], Encoders.scalaInt))) encodeDecodeTest( (3, PrimitiveData(1, 1, 1, 1, 1, 1, true)), "tuple with product encoder and flat encoder")( - ExpressionEncoder.tuple(ExpressionEncoder[Int](), ExpressionEncoder[PrimitiveData]())) + encoderFor(Encoders.tuple(Encoders.scalaInt, Encoders.product[PrimitiveData]))) encodeDecodeTest( (1, (10, 100L)), "nested tuple encoder") { - val intEnc = ExpressionEncoder[Int]() - val longEnc = ExpressionEncoder[Long]() - ExpressionEncoder.tuple(intEnc, ExpressionEncoder.tuple(intEnc, longEnc)) + val intEnc = Encoders.scalaInt + val longEnc = Encoders.scalaLong + encoderFor(Encoders.tuple(intEnc, Encoders.tuple(intEnc, longEnc))) } // test for value classes @@ -468,9 +468,8 @@ class ExpressionEncoderSuite extends CodegenInterpretedPlanTest with AnalysisTes // test for tupled encoders { - val schema = ExpressionEncoder.tuple( - ExpressionEncoder[Int](), - ExpressionEncoder[(String, Int)]()).schema + val encoder = encoderFor(Encoders.tuple(Encoders.scalaInt, Encoders.product[(String, Int)])) + val schema = encoder.schema assert(schema(0).nullable === false) assert(schema(1).nullable) assert(schema(1).dataType.asInstanceOf[StructType](0).nullable) @@ -513,11 +512,11 @@ class ExpressionEncoderSuite extends CodegenInterpretedPlanTest with AnalysisTes } test("throw exception for tuples with more than 22 elements") { - val encoders = (0 to 22).map(_ => Encoders.scalaInt.asInstanceOf[ExpressionEncoder[_]]) + val encoders = (0 to 22).map(_ => Encoders.scalaInt) checkError( exception = intercept[SparkUnsupportedOperationException] { - ExpressionEncoder.tuple(encoders) + Encoders.tupleEncoder(encoders: _*) }, condition = "_LEGACY_ERROR_TEMP_2150", parameters = Map.empty) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala index bb6d52308c192..33c9edb1cd21a 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala @@ -47,7 +47,7 @@ import org.apache.spark.sql.{Dataset, Encoders, ForeachWriter, Observation, Rela import org.apache.spark.sql.avro.{AvroDataToCatalyst, CatalystDataToAvro} import org.apache.spark.sql.catalyst.{expressions, AliasIdentifier, FunctionIdentifier, QueryPlanningTracker} import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, GlobalTempView, LocalTempView, MultiAlias, NameParameterizedQuery, PosParameterizedQuery, UnresolvedAlias, UnresolvedAttribute, UnresolvedDataFrameStar, UnresolvedDeserializer, UnresolvedExtractValue, UnresolvedFunction, UnresolvedRegex, UnresolvedRelation, UnresolvedStar, UnresolvedTranspose} -import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, ExpressionEncoder, RowEncoder} +import org.apache.spark.sql.catalyst.encoders.{encoderFor, AgnosticEncoder, ExpressionEncoder, RowEncoder} import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.UnboundRowEncoder import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression @@ -2318,16 +2318,17 @@ class SparkConnectPlanner( if (fun.getArgumentsCount != 1) { throw InvalidPlanInput("reduce requires single child expression") } - val udf = fun.getArgumentsList.asScala.map(transformExpression) match { - case collection.Seq(f: ScalaUDF) => - f + val udf = fun.getArgumentsList.asScala match { + case collection.Seq(e) + if e.hasCommonInlineUserDefinedFunction && + e.getCommonInlineUserDefinedFunction.hasScalarScalaUdf => + unpackUdf(e.getCommonInlineUserDefinedFunction) case other => throw InvalidPlanInput(s"reduce should carry a scalar scala udf, but got $other") } - assert(udf.outputEncoder.isDefined) - val tEncoder = udf.outputEncoder.get // (T, T) => T - val reduce = ReduceAggregator(udf.function)(tEncoder).toColumn.expr - TypedAggUtils.withInputType(reduce, tEncoder, dataAttributes) + val encoder = udf.outputEncoder + val reduce = ReduceAggregator(udf.function)(encoder).toColumn.expr + TypedAggUtils.withInputType(reduce, encoderFor(encoder), dataAttributes) } private def transformExpressionWithTypedReduceExpression( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 6e5dcc24e29dd..c147b6a56e024 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -42,6 +42,7 @@ import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, Query import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.catalog.HiveTableRelation import org.apache.spark.sql.catalyst.encoders._ +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, ProductEncoder, StructEncoder} import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.json.{JacksonGenerator, JSONOptions} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserUtils} @@ -78,13 +79,14 @@ private[sql] object Dataset { val DATASET_ID_TAG = TreeNodeTag[HashSet[Long]]("dataset_id") def apply[T: Encoder](sparkSession: SparkSession, logicalPlan: LogicalPlan): Dataset[T] = { - val dataset = new Dataset(sparkSession, logicalPlan, implicitly[Encoder[T]]) + val encoder = implicitly[Encoder[T]] + val dataset = new Dataset(sparkSession, logicalPlan, encoder) // Eagerly bind the encoder so we verify that the encoder matches the underlying // schema. The user will get an error if this is not the case. // optimization: it is guaranteed that [[InternalRow]] can be converted to [[Row]] so // do not do this check in that case. this check can be expensive since it requires running // the whole [[Analyzer]] to resolve the deserializer - if (dataset.exprEnc.clsTag.runtimeClass != classOf[Row]) { + if (dataset.encoder.clsTag.runtimeClass != classOf[Row]) { dataset.resolvedEnc } dataset @@ -94,7 +96,7 @@ private[sql] object Dataset { sparkSession.withActive { val qe = sparkSession.sessionState.executePlan(logicalPlan) qe.assertAnalyzed() - new Dataset[Row](qe, ExpressionEncoder(qe.analyzed.schema)) + new Dataset[Row](qe, RowEncoder.encoderFor(qe.analyzed.schema)) } def ofRows( @@ -105,7 +107,7 @@ private[sql] object Dataset { val qe = new QueryExecution( sparkSession, logicalPlan, shuffleCleanupMode = shuffleCleanupMode) qe.assertAnalyzed() - new Dataset[Row](qe, ExpressionEncoder(qe.analyzed.schema)) + new Dataset[Row](qe, RowEncoder.encoderFor(qe.analyzed.schema)) } /** A variant of ofRows that allows passing in a tracker so we can track query parsing time. */ @@ -118,7 +120,7 @@ private[sql] object Dataset { val qe = new QueryExecution( sparkSession, logicalPlan, tracker, shuffleCleanupMode = shuffleCleanupMode) qe.assertAnalyzed() - new Dataset[Row](qe, ExpressionEncoder(qe.analyzed.schema)) + new Dataset[Row](qe, RowEncoder.encoderFor(qe.analyzed.schema)) } } @@ -252,12 +254,17 @@ class Dataset[T] private[sql]( } /** - * Currently [[ExpressionEncoder]] is the only implementation of [[Encoder]], here we turn the - * passed in encoder to [[ExpressionEncoder]] explicitly, and mark it implicit so that we can use - * it when constructing new Dataset objects that have the same object type (that will be - * possibly resolved to a different schema). + * Expose the encoder as implicit so it can be used to construct new Dataset objects that have + * the same external type. */ - private[sql] implicit val exprEnc: ExpressionEncoder[T] = encoderFor(encoder) + private implicit def encoderImpl: Encoder[T] = encoder + + /** + * The actual [[ExpressionEncoder]] used by the dataset. This and its resolved counterpart should + * only be used for actual (de)serialization, the binding of Aggregator inputs, and in the rare + * cases where a plan needs to be constructed with an ExpressionEncoder. + */ + private[sql] lazy val exprEnc: ExpressionEncoder[T] = encoderFor(encoder) // The resolved `ExpressionEncoder` which can be used to turn rows to objects of type T, after // collecting rows to the driver side. @@ -265,7 +272,7 @@ class Dataset[T] private[sql]( exprEnc.resolveAndBind(logicalPlan.output, sparkSession.sessionState.analyzer) } - private implicit def classTag: ClassTag[T] = exprEnc.clsTag + private implicit def classTag: ClassTag[T] = encoder.clsTag // sqlContext must be val because a stable identifier is expected when you import implicits @transient lazy val sqlContext: SQLContext = sparkSession.sqlContext @@ -476,7 +483,7 @@ class Dataset[T] private[sql]( /** @inheritdoc */ // This is declared with parentheses to prevent the Scala compiler from treating // `ds.toDF("1")` as invoking this toDF and then apply on the returned DataFrame. - def toDF(): DataFrame = new Dataset[Row](queryExecution, ExpressionEncoder(schema)) + def toDF(): DataFrame = new Dataset[Row](queryExecution, RowEncoder.encoderFor(schema)) /** @inheritdoc */ def as[U : Encoder]: Dataset[U] = Dataset[U](sparkSession, logicalPlan) @@ -671,17 +678,17 @@ class Dataset[T] private[sql]( Some(condition.expr), JoinHint.NONE)).analyzed.asInstanceOf[Join] - implicit val tuple2Encoder: Encoder[(T, U)] = - ExpressionEncoder - .tuple(Seq(this.exprEnc, other.exprEnc), useNullSafeDeserializer = true) - .asInstanceOf[Encoder[(T, U)]] - - withTypedPlan(JoinWith.typedJoinWith( + val leftEncoder = agnosticEncoderFor(encoder) + val rightEncoder = agnosticEncoderFor(other.encoder) + val joinEncoder = ProductEncoder.tuple(Seq(leftEncoder, rightEncoder), elementsCanBeNull = true) + .asInstanceOf[Encoder[(T, U)]] + val joinWith = JoinWith.typedJoinWith( joined, sparkSession.sessionState.conf.dataFrameSelfJoinAutoResolveAmbiguity, sparkSession.sessionState.analyzer.resolver, - this.exprEnc.isSerializedAsStructForTopLevel, - other.exprEnc.isSerializedAsStructForTopLevel)) + leftEncoder.isStruct, + rightEncoder.isStruct) + new Dataset(sparkSession, joinWith, joinEncoder) } // TODO(SPARK-22947): Fix the DataFrame API. @@ -826,24 +833,29 @@ class Dataset[T] private[sql]( /** @inheritdoc */ def select[U1](c1: TypedColumn[T, U1]): Dataset[U1] = { - implicit val encoder: ExpressionEncoder[U1] = encoderFor(c1.encoder) + val encoder = agnosticEncoderFor(c1.encoder) val tc1 = withInputType(c1.named, exprEnc, logicalPlan.output) val project = Project(tc1 :: Nil, logicalPlan) - if (!encoder.isSerializedAsStructForTopLevel) { - new Dataset[U1](sparkSession, project, encoder) - } else { - // Flattens inner fields of U1 - new Dataset[Tuple1[U1]](sparkSession, project, ExpressionEncoder.tuple(encoder)).map(_._1) + val plan = encoder match { + case se: StructEncoder[U1] => + // Flatten the result. + val attribute = GetColumnByOrdinal(0, se.dataType) + val projectList = se.fields.zipWithIndex.map { + case (field, index) => + Alias(GetStructField(attribute, index, None), field.name)() + } + Project(projectList, project) + case _ => project } + new Dataset[U1](sparkSession, plan, encoder) } /** @inheritdoc */ protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = { - val encoders = columns.map(c => encoderFor(c.encoder)) + val encoders = columns.map(c => agnosticEncoderFor(c.encoder)) val namedColumns = columns.map(c => withInputType(c.named, exprEnc, logicalPlan.output)) - val execution = new QueryExecution(sparkSession, Project(namedColumns, logicalPlan)) - new Dataset(execution, ExpressionEncoder.tuple(encoders)) + new Dataset(sparkSession, Project(namedColumns, logicalPlan), ProductEncoder.tuple(encoders)) } /** @inheritdoc */ @@ -912,8 +924,8 @@ class Dataset[T] private[sql]( val executed = sparkSession.sessionState.executePlan(withGroupingKey) new KeyValueGroupedDataset( - encoderFor[K], - encoderFor[T], + implicitly[Encoder[K]], + encoder, executed, logicalPlan.output, withGroupingKey.newColumns) @@ -1387,7 +1399,11 @@ class Dataset[T] private[sql]( packageNames: Array[Byte], broadcastVars: Array[Broadcast[Object]], schema: StructType): DataFrame = { - val rowEncoder = encoder.asInstanceOf[ExpressionEncoder[Row]] + val rowEncoder: ExpressionEncoder[Row] = if (isUnTyped) { + exprEnc.asInstanceOf[ExpressionEncoder[Row]] + } else { + ExpressionEncoder(schema) + } Dataset.ofRows( sparkSession, MapPartitionsInR(func, packageNames, broadcastVars, schema, rowEncoder, logicalPlan)) @@ -2237,7 +2253,7 @@ class Dataset[T] private[sql]( /** A convenient function to wrap a set based logical plan and produce a Dataset. */ @inline private def withSetOperator[U : Encoder](logicalPlan: LogicalPlan): Dataset[U] = { - if (classTag.runtimeClass.isAssignableFrom(classOf[Row])) { + if (isUnTyped) { // Set operators widen types (change the schema), so we cannot reuse the row encoder. Dataset.ofRows(sparkSession, logicalPlan).asInstanceOf[Dataset[U]] } else { @@ -2245,6 +2261,8 @@ class Dataset[T] private[sql]( } } + private def isUnTyped: Boolean = classTag.runtimeClass.isAssignableFrom(classOf[Row]) + /** Returns a optimized plan for CommandResult, convert to `LocalRelation`. */ private def commandResultOptimized: Dataset[T] = { logicalPlan match { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala index 1ebdd57f1962b..fcad1b721eaca 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala @@ -19,7 +19,8 @@ package org.apache.spark.sql import org.apache.spark.api.java.function._ import org.apache.spark.sql.catalyst.analysis.{EliminateEventTimeWatermark, UnresolvedAttribute} -import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder} +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, ProductEncoder} +import org.apache.spark.sql.catalyst.encoders.encoderFor import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.execution.QueryExecution @@ -43,9 +44,8 @@ class KeyValueGroupedDataset[K, V] private[sql]( extends api.KeyValueGroupedDataset[K, V, Dataset] { type KVDS[KY, VL] = KeyValueGroupedDataset[KY, VL] - // Similar to [[Dataset]], we turn the passed in encoder to `ExpressionEncoder` explicitly. - private implicit val kExprEnc: ExpressionEncoder[K] = encoderFor(kEncoder) - private implicit val vExprEnc: ExpressionEncoder[V] = encoderFor(vEncoder) + private implicit def kEncoderImpl: Encoder[K] = kEncoder + private implicit def vEncoderImpl: Encoder[V] = vEncoder private def logicalPlan = queryExecution.analyzed private def sparkSession = queryExecution.sparkSession @@ -54,8 +54,8 @@ class KeyValueGroupedDataset[K, V] private[sql]( /** @inheritdoc */ def keyAs[L : Encoder]: KeyValueGroupedDataset[L, V] = new KeyValueGroupedDataset( - encoderFor[L], - vExprEnc, + implicitly[Encoder[L]], + vEncoder, queryExecution, dataAttributes, groupingAttributes) @@ -67,8 +67,8 @@ class KeyValueGroupedDataset[K, V] private[sql]( val executed = sparkSession.sessionState.executePlan(projected) new KeyValueGroupedDataset( - encoderFor[K], - encoderFor[W], + kEncoder, + implicitly[Encoder[W]], executed, withNewData.newColumns, groupingAttributes) @@ -297,20 +297,21 @@ class KeyValueGroupedDataset[K, V] private[sql]( /** @inheritdoc */ def reduceGroups(f: (V, V) => V): Dataset[(K, V)] = { - val vEncoder = encoderFor[V] val aggregator: TypedColumn[V, V] = new ReduceAggregator[V](f)(vEncoder).toColumn agg(aggregator) } /** @inheritdoc */ protected def aggUntyped(columns: TypedColumn[_, _]*): Dataset[_] = { - val encoders = columns.map(c => encoderFor(c.encoder)) - val namedColumns = columns.map(c => withInputType(c.named, vExprEnc, dataAttributes)) - val keyColumn = aggKeyColumn(kExprEnc, groupingAttributes) + val keyAgEncoder = agnosticEncoderFor(kEncoder) + val valueExprEncoder = encoderFor(vEncoder) + val encoders = columns.map(c => agnosticEncoderFor(c.encoder)) + val namedColumns = columns.map { c => + withInputType(c.named, valueExprEncoder, dataAttributes) + } + val keyColumn = aggKeyColumn(keyAgEncoder, groupingAttributes) val aggregate = Aggregate(groupingAttributes, keyColumn +: namedColumns, logicalPlan) - val execution = new QueryExecution(sparkSession, aggregate) - - new Dataset(execution, ExpressionEncoder.tuple(kExprEnc +: encoders)) + new Dataset(sparkSession, aggregate, ProductEncoder.tuple(keyAgEncoder +: encoders)) } /** @inheritdoc */ @@ -319,7 +320,7 @@ class KeyValueGroupedDataset[K, V] private[sql]( thisSortExprs: Column*)( otherSortExprs: Column*)( f: (K, Iterator[V], Iterator[U]) => IterableOnce[R]): Dataset[R] = { - implicit val uEncoder = other.vExprEnc + implicit val uEncoder = other.vEncoderImpl Dataset[R]( sparkSession, CoGroup( @@ -336,10 +337,10 @@ class KeyValueGroupedDataset[K, V] private[sql]( override def toString: String = { val builder = new StringBuilder - val kFields = kExprEnc.schema.map { f => + val kFields = kEncoder.schema.map { f => s"${f.name}: ${f.dataType.simpleString(2)}" } - val vFields = vExprEnc.schema.map { f => + val vFields = vEncoder.schema.map { f => s"${f.name}: ${f.dataType.simpleString(2)}" } builder.append("KeyValueGroupedDataset: [key: [") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala index 4e4454018e818..da4609135fd63 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala @@ -22,7 +22,6 @@ import org.apache.spark.annotation.Stable import org.apache.spark.api.python.PythonEvalType import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias -import org.apache.spark.sql.catalyst.encoders.encoderFor import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.plans.logical._ @@ -119,15 +118,12 @@ class RelationalGroupedDataset protected[sql]( /** @inheritdoc */ def as[K: Encoder, T: Encoder]: KeyValueGroupedDataset[K, T] = { - val keyEncoder = encoderFor[K] - val valueEncoder = encoderFor[T] - val (qe, groupingAttributes) = handleGroupingExpression(df.logicalPlan, df.sparkSession, groupingExprs) new KeyValueGroupedDataset( - keyEncoder, - valueEncoder, + implicitly[Encoder[K]], + implicitly[Encoder[T]], qe, df.logicalPlan.output, groupingAttributes) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala index 3832d73044078..09d9915022a65 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala @@ -504,7 +504,7 @@ case class ScalaAggregator[IN, BUF, OUT]( private[this] lazy val inputDeserializer = inputEncoder.createDeserializer() private[this] lazy val bufferSerializer = bufferEncoder.createSerializer() private[this] lazy val bufferDeserializer = bufferEncoder.createDeserializer() - private[this] lazy val outputEncoder = agg.outputEncoder.asInstanceOf[ExpressionEncoder[OUT]] + private[this] lazy val outputEncoder = encoderFor(agg.outputEncoder) private[this] lazy val outputSerializer = outputEncoder.createSerializer() def dataType: DataType = outputEncoder.objSerializer.dataType diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousTextSocketSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousTextSocketSource.scala index 420c3e3be16d6..273ffa6aefb7b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousTextSocketSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousTextSocketSource.scala @@ -32,8 +32,9 @@ import org.apache.spark.SparkEnv import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys.{HOST, PORT} import org.apache.spark.rpc.RpcEndpointRef +import org.apache.spark.sql.Encoders import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.encoderFor import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow} import org.apache.spark.sql.connector.read.InputPartition import org.apache.spark.sql.connector.read.streaming.{ContinuousPartitionReader, ContinuousPartitionReaderFactory, ContinuousStream, Offset, PartitionOffset} @@ -57,8 +58,7 @@ class TextSocketContinuousStream( implicit val defaultFormats: DefaultFormats = DefaultFormats - private val encoder = ExpressionEncoder.tuple(ExpressionEncoder[String](), - ExpressionEncoder[Timestamp]()) + private val encoder = encoderFor(Encoders.tuple(Encoders.STRING, Encoders.TIMESTAMP)) @GuardedBy("this") private var socket: Socket = _ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala index fd3df372a2d56..192b5bf65c4c5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala @@ -19,7 +19,8 @@ package org.apache.spark.sql.expressions import org.apache.spark.SparkException import org.apache.spark.sql.Encoder -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{PrimitiveBooleanEncoder, ProductEncoder} /** * An aggregator that uses a single associative and commutative reduce function. This reduce @@ -46,10 +47,10 @@ private[sql] class ReduceAggregator[T: Encoder](func: (T, T) => T) override def zero: (Boolean, T) = (false, _zero.asInstanceOf[T]) - override def bufferEncoder: Encoder[(Boolean, T)] = - ExpressionEncoder.tuple( - ExpressionEncoder[Boolean](), - encoder.asInstanceOf[ExpressionEncoder[T]]) + override def bufferEncoder: Encoder[(Boolean, T)] = { + ProductEncoder.tuple(Seq(PrimitiveBooleanEncoder, encoder.asInstanceOf[AgnosticEncoder[T]])) + .asInstanceOf[Encoder[(Boolean, T)]] + } override def outputEncoder: Encoder[T] = encoder diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/TypedAggUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/TypedAggUtils.scala index b6340a35e7703..23ceb8135fa8a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/TypedAggUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/TypedAggUtils.scala @@ -17,7 +17,9 @@ package org.apache.spark.sql.internal +import org.apache.spark.sql.Encoder import org.apache.spark.sql.catalyst.analysis.UnresolvedDeserializer +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.agnosticEncoderFor import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression @@ -25,10 +27,10 @@ import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression private[sql] object TypedAggUtils { def aggKeyColumn[A]( - encoder: ExpressionEncoder[A], + encoder: Encoder[A], groupingAttributes: Seq[Attribute]): NamedExpression = { - if (!encoder.isSerializedAsStructForTopLevel) { - assert(groupingAttributes.length == 1) + val agnosticEncoder = agnosticEncoderFor(encoder) + if (!agnosticEncoder.isStruct) { if (SQLConf.get.nameNonStructGroupingKeyAsValue) { groupingAttributes.head } else { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/FlatMapGroupsWithStateExecHelperSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/FlatMapGroupsWithStateExecHelperSuite.scala index ea6fd8ab312c9..2456999b4382a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/FlatMapGroupsWithStateExecHelperSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/FlatMapGroupsWithStateExecHelperSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.streaming.state import java.util.concurrent.atomic.AtomicInteger import org.apache.spark.sql.Encoder -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder} import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow} import org.apache.spark.sql.execution.streaming.GroupStateImpl._ import org.apache.spark.sql.streaming.StreamTest @@ -201,7 +201,7 @@ class FlatMapGroupsWithStateExecHelperSuite extends StreamTest { private def newStateManager[T: Encoder](version: Int, withTimestamp: Boolean): StateManager = { FlatMapGroupsWithStateExecHelper.createStateManager( - implicitly[Encoder[T]].asInstanceOf[ExpressionEncoder[Any]], + encoderFor[T].asInstanceOf[ExpressionEncoder[Any]], withTimestamp, version) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala index add12f7e15352..e9300464af8dc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala @@ -22,7 +22,7 @@ import java.util.UUID import org.apache.spark.{SparkIllegalArgumentException, SparkUnsupportedOperationException} import org.apache.spark.sql.Encoders -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder} import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, ListStateImplWithTTL, StatefulProcessorHandleImpl} import org.apache.spark.sql.streaming.{ListState, TimeMode, TTLConfig, ValueState} @@ -38,7 +38,7 @@ class ListStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val listState: ListState[Long] = handle.getListState[Long]("listState", Encoders.scalaLong) @@ -71,7 +71,7 @@ class ListStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ListState[Long] = handle.getListState[Long]("testState", Encoders.scalaLong) ImplicitGroupingKeyTracker.setImplicitKey("test_key") @@ -99,7 +99,7 @@ class ListStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState1: ListState[Long] = handle.getListState[Long]("testState1", Encoders.scalaLong) val testState2: ListState[Long] = handle.getListState[Long]("testState2", Encoders.scalaLong) @@ -137,7 +137,7 @@ class ListStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val listState1: ListState[Long] = handle.getListState[Long]("listState1", Encoders.scalaLong) val listState2: ListState[Long] = handle.getListState[Long]("listState2", Encoders.scalaLong) @@ -167,7 +167,7 @@ class ListStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) @@ -187,7 +187,7 @@ class ListStateSuite extends StateVariableSuiteBase { // increment batchProcessingTime, or watermark and ensure expired value is not returned val nextBatchHandle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(ttlExpirationMs)) val nextBatchTestState: ListStateImplWithTTL[String] = @@ -223,7 +223,7 @@ class ListStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val batchTimestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(batchTimestampMs)) Seq(null, Duration.ZERO, Duration.ofMinutes(-1)).foreach { ttlDuration => @@ -250,7 +250,7 @@ class ListStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.bean(classOf[POJOTestClass]).asInstanceOf[ExpressionEncoder[Any]], + encoderFor(Encoders.bean(classOf[POJOTestClass])).asInstanceOf[ExpressionEncoder[Any]], TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala index 9c322b201da8c..b067d589de904 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/MapStateSuite.scala @@ -22,7 +22,6 @@ import java.util.UUID import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.sql.Encoders -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, MapStateImplWithTTL, StatefulProcessorHandleImpl} import org.apache.spark.sql.streaming.{ListState, MapState, TimeMode, TTLConfig, ValueState} import org.apache.spark.sql.types.{BinaryType, StructType} @@ -41,7 +40,7 @@ class MapStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: MapState[String, Double] = handle.getMapState[String, Double]("testState", Encoders.STRING, Encoders.scalaDouble) @@ -75,7 +74,7 @@ class MapStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState1: MapState[Long, Double] = handle.getMapState[Long, Double]("testState1", Encoders.scalaLong, Encoders.scalaDouble) @@ -114,7 +113,7 @@ class MapStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val mapTestState1: MapState[String, Int] = handle.getMapState[String, Int]("mapTestState1", Encoders.STRING, Encoders.scalaInt) @@ -175,7 +174,7 @@ class MapStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.ProcessingTime(), + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) @@ -196,7 +195,7 @@ class MapStateSuite extends StateVariableSuiteBase { // increment batchProcessingTime, or watermark and ensure expired value is not returned val nextBatchHandle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(ttlExpirationMs)) val nextBatchTestState: MapStateImplWithTTL[String, String] = @@ -233,7 +232,7 @@ class MapStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val batchTimestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(batchTimestampMs)) Seq(null, Duration.ZERO, Duration.ofMinutes(-1)).foreach { ttlDuration => @@ -261,7 +260,7 @@ class MapStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.ProcessingTime(), + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala index e2940497e911e..48a6fd836a462 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala @@ -22,7 +22,6 @@ import java.util.UUID import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.sql.Encoders -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, StatefulProcessorHandleImpl, StatefulProcessorHandleState} import org.apache.spark.sql.streaming.{TimeMode, TTLConfig} @@ -33,9 +32,6 @@ import org.apache.spark.sql.streaming.{TimeMode, TTLConfig} */ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { - private def keyExprEncoder: ExpressionEncoder[Any] = - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]] - private def getTimeMode(timeMode: String): TimeMode = { timeMode match { case "None" => TimeMode.None() @@ -50,7 +46,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, getTimeMode(timeMode)) + UUID.randomUUID(), stringEncoder, getTimeMode(timeMode)) assert(handle.getHandleState === StatefulProcessorHandleState.CREATED) handle.getValueState[Long]("testState", Encoders.scalaLong) } @@ -91,7 +87,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, getTimeMode(timeMode)) + UUID.randomUUID(), stringEncoder, getTimeMode(timeMode)) Seq(StatefulProcessorHandleState.INITIALIZED, StatefulProcessorHandleState.DATA_PROCESSED, @@ -109,7 +105,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, TimeMode.None()) + UUID.randomUUID(), stringEncoder, TimeMode.None()) val ex = intercept[SparkUnsupportedOperationException] { handle.registerTimer(10000L) } @@ -145,7 +141,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, getTimeMode(timeMode)) + UUID.randomUUID(), stringEncoder, getTimeMode(timeMode)) handle.setHandleState(StatefulProcessorHandleState.INITIALIZED) assert(handle.getHandleState === StatefulProcessorHandleState.INITIALIZED) @@ -166,7 +162,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, getTimeMode(timeMode)) + UUID.randomUUID(), stringEncoder, getTimeMode(timeMode)) handle.setHandleState(StatefulProcessorHandleState.DATA_PROCESSED) assert(handle.getHandleState === StatefulProcessorHandleState.DATA_PROCESSED) @@ -206,7 +202,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, getTimeMode(timeMode)) + UUID.randomUUID(), stringEncoder, getTimeMode(timeMode)) Seq(StatefulProcessorHandleState.CREATED, StatefulProcessorHandleState.TIMER_PROCESSED, @@ -223,7 +219,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, TimeMode.ProcessingTime(), + UUID.randomUUID(), stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(10)) val valueStateWithTTL = handle.getValueState("testState", @@ -241,7 +237,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, TimeMode.ProcessingTime(), + UUID.randomUUID(), stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(10)) val listStateWithTTL = handle.getListState("testState", @@ -259,7 +255,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, TimeMode.ProcessingTime(), + UUID.randomUUID(), stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(10)) val mapStateWithTTL = handle.getMapState("testState", @@ -277,7 +273,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), keyExprEncoder, TimeMode.None()) + UUID.randomUUID(), stringEncoder, TimeMode.None()) handle.getValueState("testValueState", Encoders.STRING) handle.getListState("testListState", Encoders.STRING) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala index df6a3fd7b23e5..24a120be9d9af 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.streaming.state import org.apache.spark.sql.Encoders -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder} import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, TimerStateImpl} import org.apache.spark.sql.streaming.TimeMode @@ -45,7 +45,7 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key") val timerState = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) timerState.registerTimer(1L * 1000) assert(timerState.listTimers().toSet === Set(1000L)) assert(timerState.getExpiredTimers(Long.MaxValue).toSeq === Seq(("test_key", 1000L))) @@ -64,9 +64,9 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key") val timerState1 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) val timerState2 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) timerState1.registerTimer(1L * 1000) timerState2.registerTimer(15L * 1000) assert(timerState1.listTimers().toSet === Set(15000L, 1000L)) @@ -89,7 +89,7 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key1") val timerState1 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) timerState1.registerTimer(1L * 1000) timerState1.registerTimer(2L * 1000) assert(timerState1.listTimers().toSet === Set(1000L, 2000L)) @@ -97,7 +97,7 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key2") val timerState2 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) timerState2.registerTimer(15L * 1000) ImplicitGroupingKeyTracker.removeImplicitKey() @@ -122,7 +122,7 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key") val timerState = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) val timerTimerstamps = Seq(931L, 8000L, 452300L, 4200L, 90L, 1L, 2L, 8L, 3L, 35L, 6L, 9L, 5L) // register/put unordered timestamp into rocksDB timerTimerstamps.foreach(timerState.registerTimer) @@ -141,19 +141,19 @@ class TimerSuite extends StateVariableSuiteBase { ImplicitGroupingKeyTracker.setImplicitKey("test_key1") val timerState1 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) val timerTimestamps1 = Seq(64L, 32L, 1024L, 4096L, 0L, 1L) timerTimestamps1.foreach(timerState1.registerTimer) val timerState2 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) val timerTimestamps2 = Seq(931L, 8000L, 452300L, 4200L) timerTimestamps2.foreach(timerState2.registerTimer) ImplicitGroupingKeyTracker.removeImplicitKey() ImplicitGroupingKeyTracker.setImplicitKey("test_key3") val timerState3 = new TimerStateImpl(store, timeMode, - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]]) + stringEncoder) val timerTimerStamps3 = Seq(1L, 2L, 8L, 3L) timerTimerStamps3.foreach(timerState3.registerTimer) ImplicitGroupingKeyTracker.removeImplicitKey() @@ -171,7 +171,7 @@ class TimerSuite extends StateVariableSuiteBase { val store = provider.getStore(0) ImplicitGroupingKeyTracker.setImplicitKey(TestClass(1L, "k1")) val timerState = new TimerStateImpl(store, timeMode, - Encoders.product[TestClass].asInstanceOf[ExpressionEncoder[Any]]) + encoderFor(Encoders.product[TestClass]).asInstanceOf[ExpressionEncoder[Any]]) timerState.registerTimer(1L * 1000) assert(timerState.listTimers().toSet === Set(1000L)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala index 41912a4dda23b..13d758eb1b88f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala @@ -27,7 +27,7 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.{SparkException, SparkUnsupportedOperationException} import org.apache.spark.sql.Encoders -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.encoders.{encoderFor, ExpressionEncoder} import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, StatefulProcessorHandleImpl, ValueStateImplWithTTL} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.{TimeMode, TTLConfig, ValueState} @@ -49,7 +49,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val stateName = "testState" val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong) @@ -93,7 +93,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong) ImplicitGroupingKeyTracker.setImplicitKey("test_key") @@ -119,7 +119,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState1: ValueState[Long] = handle.getValueState[Long]( "testState1", Encoders.scalaLong) @@ -164,7 +164,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, - UUID.randomUUID(), Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + UUID.randomUUID(), stringEncoder, TimeMode.None()) val cfName = "$testState" val ex = intercept[SparkUnsupportedOperationException] { @@ -204,7 +204,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ValueState[Double] = handle.getValueState[Double]("testState", Encoders.scalaDouble) @@ -230,7 +230,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ValueState[Long] = handle.getValueState[Long]("testState", Encoders.scalaLong) @@ -256,7 +256,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ValueState[TestClass] = handle.getValueState[TestClass]("testState", Encoders.product[TestClass]) @@ -282,7 +282,7 @@ class ValueStateSuite extends StateVariableSuiteBase { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.None()) + stringEncoder, TimeMode.None()) val testState: ValueState[POJOTestClass] = handle.getValueState[POJOTestClass]("testState", Encoders.bean(classOf[POJOTestClass])) @@ -310,7 +310,7 @@ class ValueStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], TimeMode.ProcessingTime(), + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) @@ -330,7 +330,7 @@ class ValueStateSuite extends StateVariableSuiteBase { // increment batchProcessingTime, or watermark and ensure expired value is not returned val nextBatchHandle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(ttlExpirationMs)) val nextBatchTestState: ValueStateImplWithTTL[String] = @@ -366,7 +366,7 @@ class ValueStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val batchTimestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]], + stringEncoder, TimeMode.ProcessingTime(), batchTimestampMs = Some(batchTimestampMs)) Seq(null, Duration.ZERO, Duration.ofMinutes(-1)).foreach { ttlDuration => @@ -393,8 +393,8 @@ class ValueStateSuite extends StateVariableSuiteBase { val store = provider.getStore(0) val timestampMs = 10 val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(), - Encoders.product[TestClass].asInstanceOf[ExpressionEncoder[Any]], TimeMode.ProcessingTime(), - batchTimestampMs = Some(timestampMs)) + encoderFor(Encoders.product[TestClass]).asInstanceOf[ExpressionEncoder[Any]], + TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs)) val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1)) val testState: ValueStateImplWithTTL[POJOTestClass] = @@ -437,6 +437,8 @@ abstract class StateVariableSuiteBase extends SharedSparkSession import StateStoreTestsHelper._ + protected val stringEncoder = encoderFor(Encoders.STRING).asInstanceOf[ExpressionEncoder[Any]] + // dummy schema for initializing rocksdb provider protected def schemaForKeyRow: StructType = new StructType().add("key", BinaryType) protected def schemaForValueRow: StructType = new StructType().add("value", BinaryType) From fd8e99b9df55bf2ea29b6279a6a840ffef20ed4e Mon Sep 17 00:00:00 2001 From: Paddy Xu Date: Tue, 17 Sep 2024 23:06:05 -0400 Subject: [PATCH 051/189] [SPARK-49249][SPARK-49320] Add new tag-related APIs in Connect back to Spark Core ### What changes were proposed in this pull request? This PR adds several new tag-related APIs in Connect back to Spark Core. Following the isolation practice in the original Connect API, the newly introduced APIs also support isolation: - `interrupt{Tag,All,Operation}` can only cancel jobs created by this Spark session. - `{add,remove}Tag` and `{get,clear}Tags` only apply to jobs created by this Spark session. Instead of returning query IDs like in Spark Connect, here in Spark SQL, these methods will return SQL execution root IDs - as "query IDs" are only for Connect. ### Why are the changes needed? To close the API gap between Connect and Core. ### Does this PR introduce _any_ user-facing change? Yes, Core users can use some new APIs. ### How was this patch tested? New test added. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47815 from xupefei/reverse-api-tag. Authored-by: Paddy Xu Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/SparkSession.scala | 70 +---- .../CheckConnectJvmClientCompatibility.scala | 15 - .../scala/org/apache/spark/SparkContext.scala | 56 +++- .../apache/spark/scheduler/DAGScheduler.scala | 33 ++- .../spark/scheduler/DAGSchedulerEvent.scala | 5 +- .../apache/spark/sql/api/SparkSession.scala | 92 ++++++ .../org/apache/spark/sql/SparkSession.scala | 119 +++++++- .../spark/sql/execution/SQLExecution.scala | 205 ++++++++------ ...essionJobTaggingAndCancellationSuite.scala | 262 ++++++++++++++++++ .../sql/execution/SQLExecutionSuite.scala | 2 +- 10 files changed, 667 insertions(+), 192 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 989a7e0c174c5..aa6258a14b811 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -420,7 +420,7 @@ class SparkSession private[sql] ( * * @since 3.5.0 */ - def interruptAll(): Seq[String] = { + override def interruptAll(): Seq[String] = { client.interruptAll().getInterruptedIdsList.asScala.toSeq } @@ -433,7 +433,7 @@ class SparkSession private[sql] ( * * @since 3.5.0 */ - def interruptTag(tag: String): Seq[String] = { + override def interruptTag(tag: String): Seq[String] = { client.interruptTag(tag).getInterruptedIdsList.asScala.toSeq } @@ -446,7 +446,7 @@ class SparkSession private[sql] ( * * @since 3.5.0 */ - def interruptOperation(operationId: String): Seq[String] = { + override def interruptOperation(operationId: String): Seq[String] = { client.interruptOperation(operationId).getInterruptedIdsList.asScala.toSeq } @@ -477,65 +477,17 @@ class SparkSession private[sql] ( SparkSession.onSessionClose(this) } - /** - * Add a tag to be assigned to all the operations started by this thread in this session. - * - * Often, a unit of execution in an application consists of multiple Spark executions. - * Application programmers can use this method to group all those jobs together and give a group - * tag. The application can use `org.apache.spark.sql.SparkSession.interruptTag` to cancel all - * running running executions with this tag. For example: - * {{{ - * // In the main thread: - * spark.addTag("myjobs") - * spark.range(10).map(i => { Thread.sleep(10); i }).collect() - * - * // In a separate thread: - * spark.interruptTag("myjobs") - * }}} - * - * There may be multiple tags present at the same time, so different parts of application may - * use different tags to perform cancellation at different levels of granularity. - * - * @param tag - * The tag to be added. Cannot contain ',' (comma) character or be an empty string. - * - * @since 3.5.0 - */ - def addTag(tag: String): Unit = { - client.addTag(tag) - } + /** @inheritdoc */ + override def addTag(tag: String): Unit = client.addTag(tag) - /** - * Remove a tag previously added to be assigned to all the operations started by this thread in - * this session. Noop if such a tag was not added earlier. - * - * @param tag - * The tag to be removed. Cannot contain ',' (comma) character or be an empty string. - * - * @since 3.5.0 - */ - def removeTag(tag: String): Unit = { - client.removeTag(tag) - } + /** @inheritdoc */ + override def removeTag(tag: String): Unit = client.removeTag(tag) - /** - * Get the tags that are currently set to be assigned to all the operations started by this - * thread. - * - * @since 3.5.0 - */ - def getTags(): Set[String] = { - client.getTags() - } + /** @inheritdoc */ + override def getTags(): Set[String] = client.getTags() - /** - * Clear the current thread's operation tags. - * - * @since 3.5.0 - */ - def clearTags(): Unit = { - client.clearTags() - } + /** @inheritdoc */ + override def clearTags(): Unit = client.clearTags() /** * We cannot deserialize a connect [[SparkSession]] because of a class clash on the server side. diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala index f4043f19eb6ac..abf03cfbc6722 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala @@ -365,21 +365,6 @@ object CheckConnectJvmClientCompatibility { // Experimental ProblemFilters.exclude[DirectMissingMethodProblem]( "org.apache.spark.sql.SparkSession.registerClassFinder"), - // public - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.interruptAll"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.interruptTag"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.interruptOperation"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.addTag"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.removeTag"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.getTags"), - ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.SparkSession.clearTags"), // SparkSession#Builder ProblemFilters.exclude[DirectMissingMethodProblem]( "org.apache.spark.sql.SparkSession#Builder.remote"), diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala index 485f0abcd25ee..042179d86c31a 100644 --- a/core/src/main/scala/org/apache/spark/SparkContext.scala +++ b/core/src/main/scala/org/apache/spark/SparkContext.scala @@ -27,6 +27,7 @@ import scala.collection.Map import scala.collection.concurrent.{Map => ScalaConcurrentMap} import scala.collection.immutable import scala.collection.mutable.HashMap +import scala.concurrent.{Future, Promise} import scala.jdk.CollectionConverters._ import scala.reflect.{classTag, ClassTag} import scala.util.control.NonFatal @@ -909,10 +910,20 @@ class SparkContext(config: SparkConf) extends Logging { * * @since 3.5.0 */ - def addJobTag(tag: String): Unit = { - SparkContext.throwIfInvalidTag(tag) + def addJobTag(tag: String): Unit = addJobTags(Set(tag)) + + /** + * Add multiple tags to be assigned to all the jobs started by this thread. + * See [[addJobTag]] for more details. + * + * @param tags The tags to be added. Cannot contain ',' (comma) character. + * + * @since 4.0.0 + */ + def addJobTags(tags: Set[String]): Unit = { + tags.foreach(SparkContext.throwIfInvalidTag) val existingTags = getJobTags() - val newTags = (existingTags + tag).mkString(SparkContext.SPARK_JOB_TAGS_SEP) + val newTags = (existingTags ++ tags).mkString(SparkContext.SPARK_JOB_TAGS_SEP) setLocalProperty(SparkContext.SPARK_JOB_TAGS, newTags) } @@ -924,10 +935,20 @@ class SparkContext(config: SparkConf) extends Logging { * * @since 3.5.0 */ - def removeJobTag(tag: String): Unit = { - SparkContext.throwIfInvalidTag(tag) + def removeJobTag(tag: String): Unit = removeJobTags(Set(tag)) + + /** + * Remove multiple tags to be assigned to all the jobs started by this thread. + * See [[removeJobTag]] for more details. + * + * @param tags The tags to be removed. Cannot contain ',' (comma) character. + * + * @since 4.0.0 + */ + def removeJobTags(tags: Set[String]): Unit = { + tags.foreach(SparkContext.throwIfInvalidTag) val existingTags = getJobTags() - val newTags = (existingTags - tag).mkString(SparkContext.SPARK_JOB_TAGS_SEP) + val newTags = (existingTags -- tags).mkString(SparkContext.SPARK_JOB_TAGS_SEP) if (newTags.isEmpty) { clearJobTags() } else { @@ -2684,6 +2705,25 @@ class SparkContext(config: SparkConf) extends Logging { dagScheduler.cancelJobGroup(groupId, cancelFutureJobs = true, None) } + /** + * Cancel active jobs that have the specified tag. See `org.apache.spark.SparkContext.addJobTag`. + * + * @param tag The tag to be cancelled. Cannot contain ',' (comma) character. + * @param reason reason for cancellation. + * @return A future with [[ActiveJob]]s, allowing extraction of information such as Job ID and + * tags. + */ + private[spark] def cancelJobsWithTagWithFuture( + tag: String, + reason: String): Future[Seq[ActiveJob]] = { + SparkContext.throwIfInvalidTag(tag) + assertNotStopped() + + val cancelledJobs = Promise[Seq[ActiveJob]]() + dagScheduler.cancelJobsWithTag(tag, Some(reason), Some(cancelledJobs)) + cancelledJobs.future + } + /** * Cancel active jobs that have the specified tag. See `org.apache.spark.SparkContext.addJobTag`. * @@ -2695,7 +2735,7 @@ class SparkContext(config: SparkConf) extends Logging { def cancelJobsWithTag(tag: String, reason: String): Unit = { SparkContext.throwIfInvalidTag(tag) assertNotStopped() - dagScheduler.cancelJobsWithTag(tag, Option(reason)) + dagScheduler.cancelJobsWithTag(tag, Option(reason), cancelledJobs = None) } /** @@ -2708,7 +2748,7 @@ class SparkContext(config: SparkConf) extends Logging { def cancelJobsWithTag(tag: String): Unit = { SparkContext.throwIfInvalidTag(tag) assertNotStopped() - dagScheduler.cancelJobsWithTag(tag, None) + dagScheduler.cancelJobsWithTag(tag, reason = None, cancelledJobs = None) } /** Cancel all jobs that have been scheduled or are running. */ diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 6c824e2fdeaed..2c89fe7885d08 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -27,6 +27,7 @@ import scala.annotation.tailrec import scala.collection.Map import scala.collection.mutable import scala.collection.mutable.{HashMap, HashSet, ListBuffer} +import scala.concurrent.Promise import scala.concurrent.duration._ import scala.util.control.NonFatal @@ -1116,11 +1117,18 @@ private[spark] class DAGScheduler( /** * Cancel all jobs with a given tag. + * + * @param tag The tag to be cancelled. Cannot contain ',' (comma) character. + * @param reason reason for cancellation. + * @param cancelledJobs a promise to be completed with operation IDs being cancelled. */ - def cancelJobsWithTag(tag: String, reason: Option[String]): Unit = { + def cancelJobsWithTag( + tag: String, + reason: Option[String], + cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = { SparkContext.throwIfInvalidTag(tag) logInfo(log"Asked to cancel jobs with tag ${MDC(TAG, tag)}") - eventProcessLoop.post(JobTagCancelled(tag, reason)) + eventProcessLoop.post(JobTagCancelled(tag, reason, cancelledJobs)) } /** @@ -1234,17 +1242,22 @@ private[spark] class DAGScheduler( jobIds.foreach(handleJobCancellation(_, Option(updatedReason))) } - private[scheduler] def handleJobTagCancelled(tag: String, reason: Option[String]): Unit = { - // Cancel all jobs belonging that have this tag. + private[scheduler] def handleJobTagCancelled( + tag: String, + reason: Option[String], + cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = { + // Cancel all jobs that have all provided tags. // First finds all active jobs with this group id, and then kill stages for them. - val jobIds = activeJobs.filter { activeJob => + val jobsToBeCancelled = activeJobs.filter { activeJob => Option(activeJob.properties).exists { properties => Option(properties.getProperty(SparkContext.SPARK_JOB_TAGS)).getOrElse("") .split(SparkContext.SPARK_JOB_TAGS_SEP).filter(!_.isEmpty).toSet.contains(tag) } - }.map(_.jobId) - val updatedReason = reason.getOrElse("part of cancelled job tag %s".format(tag)) - jobIds.foreach(handleJobCancellation(_, Option(updatedReason))) + } + val updatedReason = + reason.getOrElse("part of cancelled job tags %s".format(tag)) + jobsToBeCancelled.map(_.jobId).foreach(handleJobCancellation(_, Option(updatedReason))) + cancelledJobs.map(_.success(jobsToBeCancelled.toSeq)) } private[scheduler] def handleBeginEvent(task: Task[_], taskInfo: TaskInfo): Unit = { @@ -3113,8 +3126,8 @@ private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler case JobGroupCancelled(groupId, cancelFutureJobs, reason) => dagScheduler.handleJobGroupCancelled(groupId, cancelFutureJobs, reason) - case JobTagCancelled(tag, reason) => - dagScheduler.handleJobTagCancelled(tag, reason) + case JobTagCancelled(tag, reason, cancelledJobs) => + dagScheduler.handleJobTagCancelled(tag, reason, cancelledJobs) case AllJobsCancelled => dagScheduler.doCancelAllJobs() diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala index c9ad54d1fdc7e..8932d2ef323ba 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala @@ -19,6 +19,8 @@ package org.apache.spark.scheduler import java.util.Properties +import scala.concurrent.Promise + import org.apache.spark._ import org.apache.spark.rdd.RDD import org.apache.spark.util.{AccumulatorV2, CallSite} @@ -71,7 +73,8 @@ private[scheduler] case class JobGroupCancelled( private[scheduler] case class JobTagCancelled( tagName: String, - reason: Option[String]) extends DAGSchedulerEvent + reason: Option[String], + cancelledJobs: Option[Promise[Seq[ActiveJob]]]) extends DAGSchedulerEvent private[scheduler] case object AllJobsCancelled extends DAGSchedulerEvent diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 0580931620aaa..63d4a12e11839 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -401,6 +401,98 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C @scala.annotation.varargs def addArtifacts(uri: URI*): Unit + /** + * Add a tag to be assigned to all the operations started by this thread in this session. + * + * Often, a unit of execution in an application consists of multiple Spark executions. + * Application programmers can use this method to group all those jobs together and give a group + * tag. The application can use `org.apache.spark.sql.SparkSession.interruptTag` to cancel all + * running executions with this tag. For example: + * {{{ + * // In the main thread: + * spark.addTag("myjobs") + * spark.range(10).map(i => { Thread.sleep(10); i }).collect() + * + * // In a separate thread: + * spark.interruptTag("myjobs") + * }}} + * + * There may be multiple tags present at the same time, so different parts of application may + * use different tags to perform cancellation at different levels of granularity. + * + * @param tag + * The tag to be added. Cannot contain ',' (comma) character or be an empty string. + * + * @since 4.0.0 + */ + def addTag(tag: String): Unit + + /** + * Remove a tag previously added to be assigned to all the operations started by this thread in + * this session. Noop if such a tag was not added earlier. + * + * @param tag + * The tag to be removed. Cannot contain ',' (comma) character or be an empty string. + * + * @since 4.0.0 + */ + def removeTag(tag: String): Unit + + /** + * Get the operation tags that are currently set to be assigned to all the operations started by + * this thread in this session. + * + * @since 4.0.0 + */ + def getTags(): Set[String] + + /** + * Clear the current thread's operation tags. + * + * @since 4.0.0 + */ + def clearTags(): Unit + + /** + * Request to interrupt all currently running operations of this session. + * + * @note + * This method will wait up to 60 seconds for the interruption request to be issued. + * + * @return + * Sequence of operation IDs requested to be interrupted. + * + * @since 4.0.0 + */ + def interruptAll(): Seq[String] + + /** + * Request to interrupt all currently running operations of this session with the given job tag. + * + * @note + * This method will wait up to 60 seconds for the interruption request to be issued. + * + * @return + * Sequence of operation IDs requested to be interrupted. + * + * @since 4.0.0 + */ + def interruptTag(tag: String): Seq[String] + + /** + * Request to interrupt an operation of this session, given its operation ID. + * + * @note + * This method will wait up to 60 seconds for the interruption request to be issued. + * + * @return + * The operation ID requested to be interrupted, as a single-element sequence, or an empty + * sequence if the operation is not started by this session. + * + * @since 4.0.0 + */ + def interruptOperation(operationId: String): Seq[String] + /** * Returns a [[DataFrameReader]] that can be used to read non-streaming data in as a * `DataFrame`. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index 5746b942341fc..720b77b0b9fe5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -20,8 +20,10 @@ package org.apache.spark.sql import java.net.URI import java.nio.file.Paths import java.util.{ServiceLoader, UUID} +import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ import scala.reflect.runtime.universe.TypeTag import scala.util.control.NonFatal @@ -57,7 +59,7 @@ import org.apache.spark.sql.sources.BaseRelation import org.apache.spark.sql.streaming._ import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.ExecutionListenerManager -import org.apache.spark.util.{CallSite, SparkFileUtils, Utils} +import org.apache.spark.util.{CallSite, SparkFileUtils, ThreadUtils, Utils} import org.apache.spark.util.ArrayImplicits._ /** @@ -92,7 +94,8 @@ class SparkSession private( @transient private val existingSharedState: Option[SharedState], @transient private val parentSessionState: Option[SessionState], @transient private[sql] val extensions: SparkSessionExtensions, - @transient private[sql] val initialSessionOptions: Map[String, String]) + @transient private[sql] val initialSessionOptions: Map[String, String], + @transient private val parentManagedJobTags: Map[String, String]) extends api.SparkSession[Dataset] with Logging { self => // The call site where this SparkSession was constructed. @@ -107,7 +110,12 @@ class SparkSession private( private[sql] def this( sc: SparkContext, initialSessionOptions: java.util.HashMap[String, String]) = { - this(sc, None, None, applyAndLoadExtensions(sc), initialSessionOptions.asScala.toMap) + this( + sc, + existingSharedState = None, + parentSessionState = None, + applyAndLoadExtensions(sc), initialSessionOptions.asScala.toMap, + parentManagedJobTags = Map.empty) } private[sql] def this(sc: SparkContext) = this(sc, new java.util.HashMap[String, String]()) @@ -122,6 +130,18 @@ class SparkSession private( .getOrElse(SQLConf.getFallbackConf) }) + /** Tag to mark all jobs owned by this session. */ + private[sql] lazy val sessionJobTag = s"spark-session-$sessionUUID" + + /** + * A map to hold the mapping from user-defined tags to the real tags attached to Jobs. + * Real tag have the current session ID attached: `"tag1" -> s"spark-session-$sessionUUID-tag1"`. + */ + @transient + private[sql] lazy val managedJobTags: ConcurrentHashMap[String, String] = { + new ConcurrentHashMap(parentManagedJobTags.asJava) + } + /** @inheritdoc */ def version: String = SPARK_VERSION @@ -235,7 +255,8 @@ class SparkSession private( Some(sharedState), parentSessionState = None, extensions, - initialSessionOptions) + initialSessionOptions, + parentManagedJobTags = Map.empty) } /** @@ -256,8 +277,10 @@ class SparkSession private( Some(sharedState), Some(sessionState), extensions, - Map.empty) + Map.empty, + managedJobTags.asScala.toMap) result.sessionState // force copy of SessionState + result.managedJobTags // force copy of userDefinedToRealTagsMap result } @@ -636,6 +659,83 @@ class SparkSession private( artifactManager.addLocalArtifacts(uri.flatMap(Artifact.parseArtifacts)) } + /** @inheritdoc */ + override def addTag(tag: String): Unit = { + SparkContext.throwIfInvalidTag(tag) + managedJobTags.put(tag, s"spark-session-$sessionUUID-$tag") + } + + /** @inheritdoc */ + override def removeTag(tag: String): Unit = managedJobTags.remove(tag) + + /** @inheritdoc */ + override def getTags(): Set[String] = managedJobTags.keys().asScala.toSet + + /** @inheritdoc */ + override def clearTags(): Unit = managedJobTags.clear() + + /** + * Request to interrupt all currently running SQL operations of this session. + * + * @note Only DataFrame/SQL operations started by this session can be interrupted. + * + * @note This method will wait up to 60 seconds for the interruption request to be issued. + + * @return Sequence of SQL execution IDs requested to be interrupted. + + * @since 4.0.0 + */ + override def interruptAll(): Seq[String] = + doInterruptTag(sessionJobTag, "as part of cancellation of all jobs") + + /** + * Request to interrupt all currently running SQL operations of this session with the given + * job tag. + * + * @note Only DataFrame/SQL operations started by this session can be interrupted. + * + * @note This method will wait up to 60 seconds for the interruption request to be issued. + * + * @return Sequence of SQL execution IDs requested to be interrupted. + + * @since 4.0.0 + */ + override def interruptTag(tag: String): Seq[String] = { + val realTag = managedJobTags.get(tag) + if (realTag == null) return Seq.empty + doInterruptTag(realTag, s"part of cancelled job tags $tag") + } + + private def doInterruptTag(tag: String, reason: String): Seq[String] = { + val cancelledTags = + sparkContext.cancelJobsWithTagWithFuture(tag, reason) + + ThreadUtils.awaitResult(cancelledTags, 60.seconds) + .flatMap(job => Option(job.properties.getProperty(SQLExecution.EXECUTION_ROOT_ID_KEY))) + } + + /** + * Request to interrupt a SQL operation of this session, given its SQL execution ID. + * + * @note Only DataFrame/SQL operations started by this session can be interrupted. + * + * @note This method will wait up to 60 seconds for the interruption request to be issued. + * + * @return The execution ID requested to be interrupted, as a single-element sequence, or an empty + * sequence if the operation is not started by this session. + * + * @since 4.0.0 + */ + override def interruptOperation(operationId: String): Seq[String] = { + scala.util.Try(operationId.toLong).toOption match { + case Some(executionIdToBeCancelled) => + val tagToBeCancelled = SQLExecution.executionIdJobTag(this, executionIdToBeCancelled) + doInterruptTag(tagToBeCancelled, reason = "") + case None => + throw new IllegalArgumentException("executionId must be a number in string form.") + } + } + /** @inheritdoc */ def read: DataFrameReader = new DataFrameReader(self) @@ -722,7 +822,7 @@ class SparkSession private( } /** - * Execute a block of code with the this session set as the active session, and restore the + * Execute a block of code with this session set as the active session, and restore the * previous session on completion. */ private[sql] def withActive[T](block: => T): T = { @@ -958,7 +1058,12 @@ object SparkSession extends Logging { loadExtensions(extensions) applyExtensions(sparkContext, extensions) - session = new SparkSession(sparkContext, None, None, extensions, options.toMap) + session = new SparkSession(sparkContext, + existingSharedState = None, + parentSessionState = None, + extensions, + initialSessionOptions = options.toMap, + parentManagedJobTags = Map.empty) setDefaultSession(session) setActiveSession(session) registerContextListener(sparkContext) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala index 12ff649b621e3..5db14a8662138 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala @@ -44,7 +44,7 @@ object SQLExecution extends Logging { private def nextExecutionId: Long = _nextExecutionId.getAndIncrement - private val executionIdToQueryExecution = new ConcurrentHashMap[Long, QueryExecution]() + private[sql] val executionIdToQueryExecution = new ConcurrentHashMap[Long, QueryExecution]() def getQueryExecution(executionId: Long): QueryExecution = { executionIdToQueryExecution.get(executionId) @@ -52,6 +52,9 @@ object SQLExecution extends Logging { private val testing = sys.props.contains(IS_TESTING.key) + private[sql] def executionIdJobTag(session: SparkSession, id: Long) = + s"${session.sessionJobTag}-execution-root-id-$id" + private[sql] def checkSQLExecutionId(sparkSession: SparkSession): Unit = { val sc = sparkSession.sparkContext // only throw an exception during tests. a missing execution ID should not fail a job. @@ -82,6 +85,7 @@ object SQLExecution extends Logging { // And for the root execution, rootExecutionId == executionId. if (sc.getLocalProperty(EXECUTION_ROOT_ID_KEY) == null) { sc.setLocalProperty(EXECUTION_ROOT_ID_KEY, executionId.toString) + sc.addJobTag(executionIdJobTag(sparkSession, executionId)) } val rootExecutionId = sc.getLocalProperty(EXECUTION_ROOT_ID_KEY).toLong executionIdToQueryExecution.put(executionId, queryExecution) @@ -116,92 +120,94 @@ object SQLExecution extends Logging { val redactedConfigs = sparkSession.sessionState.conf.redactOptions(modifiedConfigs) withSQLConfPropagated(sparkSession) { - var ex: Option[Throwable] = None - var isExecutedPlanAvailable = false - val startTime = System.nanoTime() - val startEvent = SparkListenerSQLExecutionStart( - executionId = executionId, - rootExecutionId = Some(rootExecutionId), - description = desc, - details = callSite.longForm, - physicalPlanDescription = "", - sparkPlanInfo = SparkPlanInfo.EMPTY, - time = System.currentTimeMillis(), - modifiedConfigs = redactedConfigs, - jobTags = sc.getJobTags(), - jobGroupId = Option(sc.getLocalProperty(SparkContext.SPARK_JOB_GROUP_ID)) - ) - try { - body match { - case Left(e) => - sc.listenerBus.post(startEvent) + withSessionTagsApplied(sparkSession) { + var ex: Option[Throwable] = None + var isExecutedPlanAvailable = false + val startTime = System.nanoTime() + val startEvent = SparkListenerSQLExecutionStart( + executionId = executionId, + rootExecutionId = Some(rootExecutionId), + description = desc, + details = callSite.longForm, + physicalPlanDescription = "", + sparkPlanInfo = SparkPlanInfo.EMPTY, + time = System.currentTimeMillis(), + modifiedConfigs = redactedConfigs, + jobTags = sc.getJobTags(), + jobGroupId = Option(sc.getLocalProperty(SparkContext.SPARK_JOB_GROUP_ID)) + ) + try { + body match { + case Left(e) => + sc.listenerBus.post(startEvent) + throw e + case Right(f) => + val planDescriptionMode = + ExplainMode.fromString(sparkSession.sessionState.conf.uiExplainMode) + val planDesc = queryExecution.explainString(planDescriptionMode) + val planInfo = try { + SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan) + } catch { + case NonFatal(e) => + logDebug("Failed to generate SparkPlanInfo", e) + // If the queryExecution already failed before this, we are not able to generate + // the the plan info, so we use and empty graphviz node to make the UI happy + SparkPlanInfo.EMPTY + } + sc.listenerBus.post( + startEvent.copy(physicalPlanDescription = planDesc, sparkPlanInfo = planInfo)) + isExecutedPlanAvailable = true + f() + } + } catch { + case e: Throwable => + ex = Some(e) throw e - case Right(f) => - val planDescriptionMode = - ExplainMode.fromString(sparkSession.sessionState.conf.uiExplainMode) - val planDesc = queryExecution.explainString(planDescriptionMode) - val planInfo = try { - SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan) - } catch { - case NonFatal(e) => - logDebug("Failed to generate SparkPlanInfo", e) - // If the queryExecution already failed before this, we are not able to generate - // the the plan info, so we use and empty graphviz node to make the UI happy - SparkPlanInfo.EMPTY - } - sc.listenerBus.post( - startEvent.copy(physicalPlanDescription = planDesc, sparkPlanInfo = planInfo)) - isExecutedPlanAvailable = true - f() - } - } catch { - case e: Throwable => - ex = Some(e) - throw e - } finally { - val endTime = System.nanoTime() - val errorMessage = ex.map { - case e: SparkThrowable => - SparkThrowableHelper.getMessage(e, ErrorMessageFormat.PRETTY) - case e => - Utils.exceptionString(e) - } - if (queryExecution.shuffleCleanupMode != DoNotCleanup - && isExecutedPlanAvailable) { - val shuffleIds = queryExecution.executedPlan match { - case ae: AdaptiveSparkPlanExec => - ae.context.shuffleIds.asScala.keys - case _ => - Iterable.empty + } finally { + val endTime = System.nanoTime() + val errorMessage = ex.map { + case e: SparkThrowable => + SparkThrowableHelper.getMessage(e, ErrorMessageFormat.PRETTY) + case e => + Utils.exceptionString(e) } - shuffleIds.foreach { shuffleId => - queryExecution.shuffleCleanupMode match { - case RemoveShuffleFiles => - // Same as what we do in ContextCleaner.doCleanupShuffle, but do not unregister - // the shuffle on MapOutputTracker, so that stage retries would be triggered. - // Set blocking to Utils.isTesting to deflake unit tests. - sc.shuffleDriverComponents.removeShuffle(shuffleId, Utils.isTesting) - case SkipMigration => - SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId) - case _ => // this should not happen + if (queryExecution.shuffleCleanupMode != DoNotCleanup + && isExecutedPlanAvailable) { + val shuffleIds = queryExecution.executedPlan match { + case ae: AdaptiveSparkPlanExec => + ae.context.shuffleIds.asScala.keys + case _ => + Iterable.empty + } + shuffleIds.foreach { shuffleId => + queryExecution.shuffleCleanupMode match { + case RemoveShuffleFiles => + // Same as what we do in ContextCleaner.doCleanupShuffle, but do not unregister + // the shuffle on MapOutputTracker, so that stage retries would be triggered. + // Set blocking to Utils.isTesting to deflake unit tests. + sc.shuffleDriverComponents.removeShuffle(shuffleId, Utils.isTesting) + case SkipMigration => + SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId) + case _ => // this should not happen + } } } + val event = SparkListenerSQLExecutionEnd( + executionId, + System.currentTimeMillis(), + // Use empty string to indicate no error, as None may mean events generated by old + // versions of Spark. + errorMessage.orElse(Some(""))) + // Currently only `Dataset.withAction` and `DataFrameWriter.runCommand` specify the + // `name` parameter. The `ExecutionListenerManager` only watches SQL executions with + // name. We can specify the execution name in more places in the future, so that + // `QueryExecutionListener` can track more cases. + event.executionName = name + event.duration = endTime - startTime + event.qe = queryExecution + event.executionFailure = ex + sc.listenerBus.post(event) } - val event = SparkListenerSQLExecutionEnd( - executionId, - System.currentTimeMillis(), - // Use empty string to indicate no error, as None may mean events generated by old - // versions of Spark. - errorMessage.orElse(Some(""))) - // Currently only `Dataset.withAction` and `DataFrameWriter.runCommand` specify the `name` - // parameter. The `ExecutionListenerManager` only watches SQL executions with name. We - // can specify the execution name in more places in the future, so that - // `QueryExecutionListener` can track more cases. - event.executionName = name - event.duration = endTime - startTime - event.qe = queryExecution - event.executionFailure = ex - sc.listenerBus.post(event) } } } finally { @@ -211,6 +217,7 @@ object SQLExecution extends Logging { // The current execution is the root execution if rootExecutionId == executionId. if (sc.getLocalProperty(EXECUTION_ROOT_ID_KEY) == executionId.toString) { sc.setLocalProperty(EXECUTION_ROOT_ID_KEY, null) + sc.removeJobTag(executionIdJobTag(sparkSession, executionId)) } sc.setLocalProperty(SPARK_JOB_INTERRUPT_ON_CANCEL, originalInterruptOnCancel) } @@ -238,15 +245,28 @@ object SQLExecution extends Logging { val sc = sparkSession.sparkContext val oldExecutionId = sc.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) withSQLConfPropagated(sparkSession) { - try { - sc.setLocalProperty(SQLExecution.EXECUTION_ID_KEY, executionId) - body - } finally { - sc.setLocalProperty(SQLExecution.EXECUTION_ID_KEY, oldExecutionId) + withSessionTagsApplied(sparkSession) { + try { + sc.setLocalProperty(SQLExecution.EXECUTION_ID_KEY, executionId) + body + } finally { + sc.setLocalProperty(SQLExecution.EXECUTION_ID_KEY, oldExecutionId) + } } } } + private[sql] def withSessionTagsApplied[T](sparkSession: SparkSession)(block: => T): T = { + val allTags = sparkSession.managedJobTags.values().asScala.toSet + sparkSession.sessionJobTag + sparkSession.sparkContext.addJobTags(allTags) + + try { + block + } finally { + sparkSession.sparkContext.removeJobTags(allTags) + } + } + /** * Wrap an action with specified SQL configs. These configs will be propagated to the executor * side via job local properties. @@ -286,10 +306,13 @@ object SQLExecution extends Logging { val originalSession = SparkSession.getActiveSession val originalLocalProps = sc.getLocalProperties SparkSession.setActiveSession(activeSession) - sc.setLocalProperties(localProps) - val res = body - // reset active session and local props. - sc.setLocalProperties(originalLocalProps) + val res = withSessionTagsApplied(activeSession) { + sc.setLocalProperties(localProps) + val res = body + // reset active session and local props. + sc.setLocalProperties(originalLocalProps) + res + } if (originalSession.nonEmpty) { SparkSession.setActiveSession(originalSession.get) } else { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala new file mode 100644 index 0000000000000..e9fd07ecf18b7 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala @@ -0,0 +1,262 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql + +import java.util.concurrent.{ConcurrentHashMap, Semaphore, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.concurrent.{ExecutionContext, Future} +import scala.jdk.CollectionConverters._ + +import org.scalatest.concurrent.Eventually +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.{LocalSparkContext, SparkContext, SparkException, SparkFunSuite} +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobEnd, SparkListenerJobStart} +import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.tags.ExtendedSQLTest +import org.apache.spark.util.ThreadUtils + +/** + * Test cases for the tagging and cancellation APIs provided by [[SparkSession]]. + */ +@ExtendedSQLTest +class SparkSessionJobTaggingAndCancellationSuite + extends SparkFunSuite + with Eventually + with LocalSparkContext { + + override def afterEach(): Unit = { + try { + // This suite should not interfere with the other test suites. + SparkSession.getActiveSession.foreach(_.stop()) + SparkSession.clearActiveSession() + SparkSession.getDefaultSession.foreach(_.stop()) + SparkSession.clearDefaultSession() + resetSparkContext() + } finally { + super.afterEach() + } + } + + test("Tags are not inherited by new sessions") { + val session = SparkSession.builder().master("local").getOrCreate() + + assert(session.getTags() == Set()) + session.addTag("one") + assert(session.getTags() == Set("one")) + + val newSession = session.newSession() + assert(newSession.getTags() == Set()) + } + + test("Tags are inherited by cloned sessions") { + val session = SparkSession.builder().master("local").getOrCreate() + + assert(session.getTags() == Set()) + session.addTag("one") + assert(session.getTags() == Set("one")) + + val clonedSession = session.cloneSession() + assert(clonedSession.getTags() == Set("one")) + clonedSession.addTag("two") + assert(clonedSession.getTags() == Set("one", "two")) + + // Tags are not propagated back to the original session + assert(session.getTags() == Set("one")) + } + + test("Tags set from session are prefixed with session UUID") { + sc = new SparkContext("local[2]", "test") + val session = SparkSession.builder().sparkContext(sc).getOrCreate() + import session.implicits._ + + val sem = new Semaphore(0) + sc.addSparkListener(new SparkListener { + override def onJobStart(jobStart: SparkListenerJobStart): Unit = { + sem.release() + } + }) + + session.addTag("one") + Future { + session.range(1, 10000).map { i => Thread.sleep(100); i }.count() + }(ExecutionContext.global) + + assert(sem.tryAcquire(1, 1, TimeUnit.MINUTES)) + val activeJobsFuture = + session.sparkContext.cancelJobsWithTagWithFuture(session.managedJobTags.get("one"), "reason") + val activeJob = ThreadUtils.awaitResult(activeJobsFuture, 60.seconds).head + val actualTags = activeJob.properties.getProperty(SparkContext.SPARK_JOB_TAGS) + .split(SparkContext.SPARK_JOB_TAGS_SEP) + assert(actualTags.toSet == Set( + session.sessionJobTag, + s"${session.sessionJobTag}-one", + SQLExecution.executionIdJobTag( + session, + activeJob.properties.get(SQLExecution.EXECUTION_ROOT_ID_KEY).asInstanceOf[String].toLong))) + } + + test("Cancellation APIs in SparkSession are isolated") { + sc = new SparkContext("local[2]", "test") + val globalSession = SparkSession.builder().sparkContext(sc).getOrCreate() + var (sessionA, sessionB, sessionC): (SparkSession, SparkSession, SparkSession) = + (null, null, null) + + // global ExecutionContext has only 2 threads in Apache Spark CI + // create own thread pool for four Futures used in this test + val numThreads = 3 + val fpool = ThreadUtils.newForkJoinPool("job-tags-test-thread-pool", numThreads) + val executionContext = ExecutionContext.fromExecutorService(fpool) + + try { + // Add a listener to release the semaphore once jobs are launched. + val sem = new Semaphore(0) + val jobEnded = new AtomicInteger(0) + val jobProperties: ConcurrentHashMap[Int, java.util.Properties] = new ConcurrentHashMap() + + sc.addSparkListener(new SparkListener { + override def onJobStart(jobStart: SparkListenerJobStart): Unit = { + jobProperties.put(jobStart.jobId, jobStart.properties) + sem.release() + } + + override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit = { + sem.release() + jobEnded.incrementAndGet() + } + }) + + // Note: since tags are added in the Future threads, they don't need to be cleared in between. + val jobA = Future { + sessionA = globalSession.cloneSession() + import globalSession.implicits._ + + assert(sessionA.getTags() == Set()) + sessionA.addTag("two") + assert(sessionA.getTags() == Set("two")) + sessionA.clearTags() // check that clearing all tags works + assert(sessionA.getTags() == Set()) + sessionA.addTag("one") + assert(sessionA.getTags() == Set("one")) + try { + sessionA.range(1, 10000).map { i => Thread.sleep(100); i }.count() + } finally { + sessionA.clearTags() // clear for the case of thread reuse by another Future + } + }(executionContext) + val jobB = Future { + sessionB = globalSession.cloneSession() + import globalSession.implicits._ + + assert(sessionB.getTags() == Set()) + sessionB.addTag("one") + sessionB.addTag("two") + sessionB.addTag("one") + sessionB.addTag("two") // duplicates shouldn't matter + assert(sessionB.getTags() == Set("one", "two")) + try { + sessionB.range(1, 10000, 2).map { i => Thread.sleep(100); i }.count() + } finally { + sessionB.clearTags() // clear for the case of thread reuse by another Future + } + }(executionContext) + val jobC = Future { + sessionC = globalSession.cloneSession() + import globalSession.implicits._ + + sessionC.addTag("foo") + sessionC.removeTag("foo") + assert(sessionC.getTags() == Set()) // check that remove works removing the last tag + sessionC.addTag("boo") + try { + sessionC.range(1, 10000, 2).map { i => Thread.sleep(100); i }.count() + } finally { + sessionC.clearTags() // clear for the case of thread reuse by another Future + } + }(executionContext) + + // Block until four jobs have started. + assert(sem.tryAcquire(3, 1, TimeUnit.MINUTES)) + + // Tags are applied + assert(jobProperties.size == 3) + for (ss <- Seq(sessionA, sessionB, sessionC)) { + val jobProperty = jobProperties.values().asScala.filter(_.get(SparkContext.SPARK_JOB_TAGS) + .asInstanceOf[String].contains(ss.sessionUUID)) + assert(jobProperty.size == 1) + val tags = jobProperty.head.get(SparkContext.SPARK_JOB_TAGS).asInstanceOf[String] + .split(SparkContext.SPARK_JOB_TAGS_SEP) + + val executionRootIdTag = SQLExecution.executionIdJobTag( + ss, + jobProperty.head.get(SQLExecution.EXECUTION_ROOT_ID_KEY).asInstanceOf[String].toLong) + val userTagsPrefix = s"spark-session-${ss.sessionUUID}-" + + ss match { + case s if s == sessionA => assert(tags.toSet == Set( + s.sessionJobTag, executionRootIdTag, s"${userTagsPrefix}one")) + case s if s == sessionB => assert(tags.toSet == Set( + s.sessionJobTag, executionRootIdTag, s"${userTagsPrefix}one", s"${userTagsPrefix}two")) + case s if s == sessionC => assert(tags.toSet == Set( + s.sessionJobTag, executionRootIdTag, s"${userTagsPrefix}boo")) + } + } + + // Global session cancels nothing + assert(globalSession.interruptAll().isEmpty) + assert(globalSession.interruptTag("one").isEmpty) + assert(globalSession.interruptTag("two").isEmpty) + for (i <- SQLExecution.executionIdToQueryExecution.keys().asScala) { + assert(globalSession.interruptOperation(i.toString).isEmpty) + } + assert(jobEnded.intValue == 0) + + // One job cancelled + for (i <- SQLExecution.executionIdToQueryExecution.keys().asScala) { + sessionC.interruptOperation(i.toString) + } + val eC = intercept[SparkException] { + ThreadUtils.awaitResult(jobC, 1.minute) + }.getCause + assert(eC.getMessage contains "cancelled") + assert(sem.tryAcquire(1, 1, TimeUnit.MINUTES)) + assert(jobEnded.intValue == 1) + + // Another job cancelled + assert(sessionA.interruptTag("one").size == 1) + val eA = intercept[SparkException] { + ThreadUtils.awaitResult(jobA, 1.minute) + }.getCause + assert(eA.getMessage contains "cancelled job tags one") + assert(sem.tryAcquire(1, 1, TimeUnit.MINUTES)) + assert(jobEnded.intValue == 2) + + // The last job cancelled + sessionB.interruptAll() + val eB = intercept[SparkException] { + ThreadUtils.awaitResult(jobB, 1.minute) + }.getCause + assert(eB.getMessage contains "cancellation of all jobs") + assert(sem.tryAcquire(1, 1, TimeUnit.MINUTES)) + assert(jobEnded.intValue == 3) + } finally { + fpool.shutdownNow() + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala index 94d33731b6de5..059a4c9b83763 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala @@ -228,7 +228,7 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper { spark.range(1).collect() spark.sparkContext.listenerBus.waitUntilEmpty() - assert(jobTags.contains(jobTag)) + assert(jobTags.get.contains(jobTag)) assert(sqlJobTags.contains(jobTag)) } finally { spark.sparkContext.removeJobTag(jobTag) From 4590538df095b20c0736ecc992ed9c0dfb926c0e Mon Sep 17 00:00:00 2001 From: panbingkun Date: Tue, 17 Sep 2024 21:14:52 -0700 Subject: [PATCH 052/189] [SPARK-49682][BUILD] Upgrade joda-time to 2.13.0 ### What changes were proposed in this pull request? The pr aims to upgrade joda-time from `2.12.7` to `2.13.0`. ### Why are the changes needed? The version `DateTimeZone` data updated to version `2024bgtz`. The full release notes: https://www.joda.org/joda-time/changes-report.html#a2.13.0 ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48130 from panbingkun/SPARK-49682. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 2 +- pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index e1ac039f25467..9871cc0bca04f 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -146,7 +146,7 @@ jjwt-api/0.12.6//jjwt-api-0.12.6.jar jline/2.14.6//jline-2.14.6.jar jline/3.25.1//jline-3.25.1.jar jna/5.14.0//jna-5.14.0.jar -joda-time/2.12.7//joda-time-2.12.7.jar +joda-time/2.13.0//joda-time-2.13.0.jar jodd-core/3.5.2//jodd-core-3.5.2.jar jpam/1.1//jpam-1.1.jar json/1.8//json-1.8.jar diff --git a/pom.xml b/pom.xml index b9f28eb619258..694ea31e6f377 100644 --- a/pom.xml +++ b/pom.xml @@ -199,7 +199,7 @@ 2.11.0 3.1.9 3.0.12 - 2.12.7 + 2.13.0 3.5.2 3.0.0 2.2.11 From 7de71a2ec78d985c2a045f13c1275101b126cec4 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Wed, 18 Sep 2024 13:54:52 +0800 Subject: [PATCH 053/189] [SPARK-49495][DOCS][FOLLOWUP] Fix Pandoc installation for GitHub Pages publication action ### What changes were proposed in this pull request? Action 'pandoc/actions/setup' is now allowed by the ASF organization account. This followup makes the installation step manual. ### Why are the changes needed? fix ci ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? https://github.com/yaooqinn/spark/actions/runs/10914663049/job/30293151174 ### Was this patch authored or co-authored using generative AI tooling? no Closes #48136 from yaooqinn/SPARK-49495-F. Authored-by: Kent Yao Signed-off-by: Kent Yao --- .github/workflows/pages.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 083620427c015..f10dadf315a1b 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -63,9 +63,9 @@ jobs: ruby-version: '3.3' bundler-cache: true - name: Install Pandoc - uses: pandoc/actions/setup@d6abb76f6c8a1a9a5e15a5190c96a02aabffd1ee - with: - version: 3.3 + run: | + sudo apt-get update -y + sudo apt-get install pandoc - name: Install dependencies for documentation generation run: | cd docs From b86e5d2ab1fb17f8dcbb5b4d50f3361494270438 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Wed, 18 Sep 2024 07:44:42 -0700 Subject: [PATCH 054/189] [SPARK-49495][DOCS][FOLLOWUP] Enable GitHub Pages settings via .asf.yml ### What changes were proposed in this pull request? A followup of SPARK-49495 to enable GitHub Pages settings via [.asf.yaml](https://cwiki.apache.org/confluence/pages/viewpage.action?spaceKey=INFRA&title=git+-+.asf.yaml+features#Git.asf.yamlfeatures-GitHubPages) ### Why are the changes needed? Meet the requirement for `actions/configure-pagesv5` action ``` Run actions/configure-pagesv5 with: token: *** enablement: false env: SPARK_TESTING: 1 RELEASE_VERSION: In-Progress JAVA_HOME: /opt/hostedtoolcache/Java_Zulu_jdk/17.0.1[2](https://github.com/apache/spark/actions/runs/10916383676/job/30297716064#step:10:2)-7/x64 JAVA_HOME_17_X64: /opt/hostedtoolcache/Java_Zulu_jdk/17.0.12-7/x64 pythonLocation: /opt/hostedtoolcache/Python/[3](https://github.com/apache/spark/actions/runs/10916383676/job/30297716064#step:10:3).9.19/x64 PKG_CONFIG_PATH: /opt/hostedtoolcache/Python/3.9.19/x6[4](https://github.com/apache/spark/actions/runs/10916383676/job/30297716064#step:10:4)/lib/pkgconfig Python_ROOT_DIR: /opt/hostedtoolcache/Python/3.9.19/x[6](https://github.com/apache/spark/actions/runs/10916383676/job/30297716064#step:10:6)4 Python2_ROOT_DIR: /opt/hostedtoolcache/Python/3.9.19/x64 Python3_ROOT_DIR: /opt/hostedtoolcache/Python/3.[9](https://github.com/apache/spark/actions/runs/10916383676/job/30297716064#step:10:9).19/x64 LD_LIBRARY_PATH: /opt/hostedtoolcache/Python/3.9.19/x64/lib Error: Get Pages site failed. Please verify that the repository has Pages enabled and configured to build using GitHub Actions, or consider exploring the `enablement` parameter for this action. Error: Not Found - https://docs.github.com/rest/pages/pages#get-a-apiname-pages-site Error: HttpError: Not Found - https://docs.github.com/rest/pages/pages#get-a-apiname-pages-site ``` ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? NA ### Was this patch authored or co-authored using generative AI tooling? no Closes #48141 from yaooqinn/SPARK-49495-FF. Authored-by: Kent Yao Signed-off-by: Dongjoon Hyun --- .asf.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.asf.yaml b/.asf.yaml index 22042b355b2fa..91a5f9b2bb1a2 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -31,6 +31,8 @@ github: merge: false squash: true rebase: true + ghp_branch: master + ghp_path: /docs/_site notifications: pullrequests: reviews@spark.apache.org From ed3a9b1aa92957015592b399167a960b68b73beb Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 18 Sep 2024 09:28:09 -0700 Subject: [PATCH 055/189] [SPARK-49691][PYTHON][CONNECT] Function `substring` should accept column names ### What changes were proposed in this pull request? Function `substring` should accept column names ### Why are the changes needed? Bug fix: ``` In [1]: >>> import pyspark.sql.functions as sf ...: >>> df = spark.createDataFrame([('Spark', 2, 3)], ['s', 'p', 'l']) ...: >>> df.select('*', sf.substring('s', 'p', 'l')).show() ``` works in PySpark Classic, but fail in Connect with: ``` NumberFormatException Traceback (most recent call last) Cell In[2], line 1 ----> 1 df.select('*', sf.substring('s', 'p', 'l')).show() File ~/Dev/spark/python/pyspark/sql/connect/dataframe.py:1170, in DataFrame.show(self, n, truncate, vertical) 1169 def show(self, n: int = 20, truncate: Union[bool, int] = True, vertical: bool = False) -> None: -> 1170 print(self._show_string(n, truncate, vertical)) File ~/Dev/spark/python/pyspark/sql/connect/dataframe.py:927, in DataFrame._show_string(self, n, truncate, vertical) 910 except ValueError: 911 raise PySparkTypeError( 912 errorClass="NOT_BOOL", 913 messageParameters={ (...) 916 }, 917 ) 919 table, _ = DataFrame( 920 plan.ShowString( 921 child=self._plan, 922 num_rows=n, 923 truncate=_truncate, 924 vertical=vertical, 925 ), 926 session=self._session, --> 927 )._to_table() 928 return table[0][0].as_py() File ~/Dev/spark/python/pyspark/sql/connect/dataframe.py:1844, in DataFrame._to_table(self) 1842 def _to_table(self) -> Tuple["pa.Table", Optional[StructType]]: 1843 query = self._plan.to_proto(self._session.client) -> 1844 table, schema, self._execution_info = self._session.client.to_table( 1845 query, self._plan.observations 1846 ) 1847 assert table is not None 1848 return (table, schema) File ~/Dev/spark/python/pyspark/sql/connect/client/core.py:892, in SparkConnectClient.to_table(self, plan, observations) 890 req = self._execute_plan_request_with_metadata() 891 req.plan.CopyFrom(plan) --> 892 table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(req, observations) 894 # Create a query execution object. 895 ei = ExecutionInfo(metrics, observed_metrics) File ~/Dev/spark/python/pyspark/sql/connect/client/core.py:1517, in SparkConnectClient._execute_and_fetch(self, req, observations, self_destruct) 1514 properties: Dict[str, Any] = {} 1516 with Progress(handlers=self._progress_handlers, operation_id=req.operation_id) as progress: -> 1517 for response in self._execute_and_fetch_as_iterator( 1518 req, observations, progress=progress 1519 ): 1520 if isinstance(response, StructType): 1521 schema = response File ~/Dev/spark/python/pyspark/sql/connect/client/core.py:1494, in SparkConnectClient._execute_and_fetch_as_iterator(self, req, observations, progress) 1492 raise kb 1493 except Exception as error: -> 1494 self._handle_error(error) File ~/Dev/spark/python/pyspark/sql/connect/client/core.py:1764, in SparkConnectClient._handle_error(self, error) 1762 self.thread_local.inside_error_handling = True 1763 if isinstance(error, grpc.RpcError): -> 1764 self._handle_rpc_error(error) 1765 elif isinstance(error, ValueError): 1766 if "Cannot invoke RPC" in str(error) and "closed" in str(error): File ~/Dev/spark/python/pyspark/sql/connect/client/core.py:1840, in SparkConnectClient._handle_rpc_error(self, rpc_error) 1837 if info.metadata["errorClass"] == "INVALID_HANDLE.SESSION_CHANGED": 1838 self._closed = True -> 1840 raise convert_exception( 1841 info, 1842 status.message, 1843 self._fetch_enriched_error(info), 1844 self._display_server_stack_trace(), 1845 ) from None 1847 raise SparkConnectGrpcException(status.message) from None 1848 else: NumberFormatException: [CAST_INVALID_INPUT] The value 'p' of the type "STRING" cannot be cast to "INT" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. SQLSTATE: 22018 ... ``` ### Does this PR introduce _any_ user-facing change? yes, Function `substring` in Connect can properly handle column names ### How was this patch tested? new doctests ### Was this patch authored or co-authored using generative AI tooling? No Closes #48135 from zhengruifeng/py_substring_fix. Authored-by: Ruifeng Zheng Signed-off-by: Dongjoon Hyun --- .../pyspark/sql/connect/functions/builtin.py | 10 ++- python/pyspark/sql/functions/builtin.py | 63 ++++++++++++++++--- 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index 031e7c22542d2..2870d9c408b6b 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -2488,8 +2488,14 @@ def sentences( sentences.__doc__ = pysparkfuncs.sentences.__doc__ -def substring(str: "ColumnOrName", pos: int, len: int) -> Column: - return _invoke_function("substring", _to_col(str), lit(pos), lit(len)) +def substring( + str: "ColumnOrName", + pos: Union["ColumnOrName", int], + len: Union["ColumnOrName", int], +) -> Column: + _pos = lit(pos) if isinstance(pos, int) else _to_col(pos) + _len = lit(len) if isinstance(len, int) else _to_col(len) + return _invoke_function("substring", _to_col(str), _pos, _len) substring.__doc__ = pysparkfuncs.substring.__doc__ diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 781bf3d9f83a2..c0730b193bc72 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -11309,7 +11309,9 @@ def sentences( @_try_remote_functions def substring( - str: "ColumnOrName", pos: Union["ColumnOrName", int], len: Union["ColumnOrName", int] + str: "ColumnOrName", + pos: Union["ColumnOrName", int], + len: Union["ColumnOrName", int], ) -> Column: """ Substring starts at `pos` and is of length `len` when str is String type or @@ -11348,16 +11350,59 @@ def substring( Examples -------- + Example 1: Using literal integers as arguments + + >>> import pyspark.sql.functions as sf >>> df = spark.createDataFrame([('abcd',)], ['s',]) - >>> df.select(substring(df.s, 1, 2).alias('s')).collect() - [Row(s='ab')] + >>> df.select('*', sf.substring(df.s, 1, 2)).show() + +----+------------------+ + | s|substring(s, 1, 2)| + +----+------------------+ + |abcd| ab| + +----+------------------+ + + Example 2: Using columns as arguments + + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([('Spark', 2, 3)], ['s', 'p', 'l']) + >>> df.select('*', sf.substring(df.s, 2, df.l)).show() + +-----+---+---+------------------+ + | s| p| l|substring(s, 2, l)| + +-----+---+---+------------------+ + |Spark| 2| 3| par| + +-----+---+---+------------------+ + + >>> df.select('*', sf.substring(df.s, df.p, 3)).show() + +-----+---+---+------------------+ + | s| p| l|substring(s, p, 3)| + +-----+---+---+------------------+ + |Spark| 2| 3| par| + +-----+---+---+------------------+ + + >>> df.select('*', sf.substring(df.s, df.p, df.l)).show() + +-----+---+---+------------------+ + | s| p| l|substring(s, p, l)| + +-----+---+---+------------------+ + |Spark| 2| 3| par| + +-----+---+---+------------------+ + + Example 3: Using column names as arguments + + >>> import pyspark.sql.functions as sf >>> df = spark.createDataFrame([('Spark', 2, 3)], ['s', 'p', 'l']) - >>> df.select(substring(df.s, 2, df.l).alias('s')).collect() - [Row(s='par')] - >>> df.select(substring(df.s, df.p, 3).alias('s')).collect() - [Row(s='par')] - >>> df.select(substring(df.s, df.p, df.l).alias('s')).collect() - [Row(s='par')] + >>> df.select('*', sf.substring(df.s, 2, 'l')).show() + +-----+---+---+------------------+ + | s| p| l|substring(s, 2, l)| + +-----+---+---+------------------+ + |Spark| 2| 3| par| + +-----+---+---+------------------+ + + >>> df.select('*', sf.substring('s', 'p', 'l')).show() + +-----+---+---+------------------+ + | s| p| l|substring(s, p, l)| + +-----+---+---+------------------+ + |Spark| 2| 3| par| + +-----+---+---+------------------+ """ pos = _enum_to_value(pos) pos = lit(pos) if isinstance(pos, int) else pos From fbf81ebaef49baa4c19a936fb3884c2e62e6a49b Mon Sep 17 00:00:00 2001 From: xuping <13289341606@163.com> Date: Wed, 18 Sep 2024 22:06:00 +0200 Subject: [PATCH 056/189] [SPARK-47263][SQL] Assign names to the legacy conditions _LEGACY_ERROR_TEMP_13[44-46] ### What changes were proposed in this pull request? rename err class _LEGACY_ERROR_TEMP_13[44-46]: 44 removed, 45 to DEFAULT_UNSUPPORTED, 46 to ADD_DEFAULT_UNSUPPORTED ### Why are the changes needed? replace legacy err class name with semantically explicits. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Re run the UT class modified in the PR (org.apache.spark.sql.sources.InsertSuite & org.apache.spark.sql.types.StructTypeSuite) ### Was this patch authored or co-authored using generative AI tooling? No Closes #46320 from PaysonXu/SPARK-47263. Authored-by: xuping <13289341606@163.com> Signed-off-by: Max Gekk --- .../resources/error/error-conditions.json | 27 +++++++++---------- .../util/ResolveDefaultColumnsUtil.scala | 9 ++++--- .../sql/errors/QueryCompilationErrors.scala | 16 +++-------- .../spark/sql/types/StructTypeSuite.scala | 23 +++++++++++----- .../spark/sql/sources/InsertSuite.scala | 6 ++--- 5 files changed, 41 insertions(+), 40 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 25dd676c4aff9..6463cc2c12da7 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1,4 +1,10 @@ { + "ADD_DEFAULT_UNSUPPORTED" : { + "message" : [ + "Failed to execute command because DEFAULT values are not supported when adding new columns to previously existing target data source with table provider: \"\"." + ], + "sqlState" : "42623" + }, "AGGREGATE_FUNCTION_WITH_NONDETERMINISTIC_EXPRESSION" : { "message" : [ "Non-deterministic expression should not appear in the arguments of an aggregate function." @@ -1096,6 +1102,12 @@ ], "sqlState" : "42608" }, + "DEFAULT_UNSUPPORTED" : { + "message" : [ + "Failed to execute command because DEFAULT values are not supported for target data source with table provider: \"\"." + ], + "sqlState" : "42623" + }, "DISTINCT_WINDOW_FUNCTION_UNSUPPORTED" : { "message" : [ "Distinct window functions are not supported: ." @@ -6673,21 +6685,6 @@ "Sinks cannot request distribution and ordering in continuous execution mode." ] }, - "_LEGACY_ERROR_TEMP_1344" : { - "message" : [ - "Invalid DEFAULT value for column : fails to parse as a valid literal value." - ] - }, - "_LEGACY_ERROR_TEMP_1345" : { - "message" : [ - "Failed to execute command because DEFAULT values are not supported for target data source with table provider: \"\"." - ] - }, - "_LEGACY_ERROR_TEMP_1346" : { - "message" : [ - "Failed to execute command because DEFAULT values are not supported when adding new columns to previously existing target data source with table provider: \"\"." - ] - }, "_LEGACY_ERROR_TEMP_2000" : { "message" : [ ". If necessary set to false to bypass this error." diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala index 8b7392e71249e..693ac8d94dbcf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.util import scala.collection.mutable.ArrayBuffer -import org.apache.spark.{SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys._ import org.apache.spark.sql.AnalysisException @@ -412,8 +412,11 @@ object ResolveDefaultColumns extends QueryErrorsBase case _: ExprLiteral | _: Cast => expr } } catch { - case _: AnalysisException | _: MatchError => - throw QueryCompilationErrors.failedToParseExistenceDefaultAsLiteral(field.name, text) + // AnalysisException thrown from analyze is already formatted, throw it directly. + case ae: AnalysisException => throw ae + case _: MatchError => + throw SparkException.internalError(s"parse existence default as literal err," + + s" field name: ${field.name}, value: $text") } // The expression should be a literal value by this point, possibly wrapped in a cast // function. This is enforced by the execution of commands that assign default values. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index f268ef85ef1dd..e324d4e9d2edb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -3516,29 +3516,21 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat "cond" -> toSQLExpr(cond))) } - def failedToParseExistenceDefaultAsLiteral(fieldName: String, defaultValue: String): Throwable = { - new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1344", - messageParameters = Map( - "fieldName" -> fieldName, - "defaultValue" -> defaultValue)) - } - def defaultReferencesNotAllowedInDataSource( statementType: String, dataSource: String): Throwable = { new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1345", + errorClass = "DEFAULT_UNSUPPORTED", messageParameters = Map( - "statementType" -> statementType, + "statementType" -> toSQLStmt(statementType), "dataSource" -> dataSource)) } def addNewDefaultColumnToExistingTableNotAllowed( statementType: String, dataSource: String): Throwable = { new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1346", + errorClass = "ADD_DEFAULT_UNSUPPORTED", messageParameters = Map( - "statementType" -> statementType, + "statementType" -> toSQLStmt(statementType), "dataSource" -> dataSource)) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala index 5ec1525bf9b61..6a67525dd02d3 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala @@ -564,7 +564,6 @@ class StructTypeSuite extends SparkFunSuite with SQLHelper { .putString(ResolveDefaultColumns.EXISTS_DEFAULT_COLUMN_METADATA_KEY, "1 + 1") .putString(ResolveDefaultColumns.CURRENT_DEFAULT_COLUMN_METADATA_KEY, "1 + 1") .build()))) - val error = "fails to parse as a valid literal value" assert(ResolveDefaultColumns.existenceDefaultValues(source2).length == 1) assert(ResolveDefaultColumns.existenceDefaultValues(source2)(0) == 2) @@ -576,9 +575,13 @@ class StructTypeSuite extends SparkFunSuite with SQLHelper { .putString(ResolveDefaultColumns.EXISTS_DEFAULT_COLUMN_METADATA_KEY, "invalid") .putString(ResolveDefaultColumns.CURRENT_DEFAULT_COLUMN_METADATA_KEY, "invalid") .build()))) - assert(intercept[AnalysisException] { - ResolveDefaultColumns.existenceDefaultValues(source3) - }.getMessage.contains(error)) + + checkError( + exception = intercept[AnalysisException]{ + ResolveDefaultColumns.existenceDefaultValues(source3) + }, + condition = "INVALID_DEFAULT_VALUE.UNRESOLVED_EXPRESSION", + parameters = Map("statement" -> "", "colName" -> "`c1`", "defaultValue" -> "invalid")) // Negative test: StructType.defaultValues fails because the existence default value fails to // resolve. @@ -592,9 +595,15 @@ class StructTypeSuite extends SparkFunSuite with SQLHelper { ResolveDefaultColumns.CURRENT_DEFAULT_COLUMN_METADATA_KEY, "(SELECT 'abc' FROM missingtable)") .build()))) - assert(intercept[AnalysisException] { - ResolveDefaultColumns.existenceDefaultValues(source4) - }.getMessage.contains(error)) + + checkError( + exception = intercept[AnalysisException]{ + ResolveDefaultColumns.existenceDefaultValues(source4) + }, + condition = "INVALID_DEFAULT_VALUE.SUBQUERY_EXPRESSION", + parameters = Map("statement" -> "", + "colName" -> "`c1`", + "defaultValue" -> "(SELECT 'abc' FROM missingtable)")) } test("SPARK-46629: Test STRUCT DDL with NOT NULL round trip") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala index 57655a58a694d..41447d8af5740 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala @@ -1998,7 +1998,7 @@ class InsertSuite extends DataSourceTest with SharedSparkSession { exception = intercept[AnalysisException] { sql(s"create table t(a string default 'abc') using parquet") }, - condition = "_LEGACY_ERROR_TEMP_1345", + condition = "DEFAULT_UNSUPPORTED", parameters = Map("statementType" -> "CREATE TABLE", "dataSource" -> "parquet")) withTable("t") { sql(s"create table t(a string, b int) using parquet") @@ -2006,7 +2006,7 @@ class InsertSuite extends DataSourceTest with SharedSparkSession { exception = intercept[AnalysisException] { sql("alter table t add column s bigint default 42") }, - condition = "_LEGACY_ERROR_TEMP_1345", + condition = "DEFAULT_UNSUPPORTED", parameters = Map( "statementType" -> "ALTER TABLE ADD COLUMNS", "dataSource" -> "parquet")) @@ -2314,7 +2314,7 @@ class InsertSuite extends DataSourceTest with SharedSparkSession { // provider is now in the denylist. sql(s"alter table t1 add column (b string default 'abc')") }, - condition = "_LEGACY_ERROR_TEMP_1346", + condition = "ADD_DEFAULT_UNSUPPORTED", parameters = Map( "statementType" -> "ALTER TABLE ADD COLUMNS", "dataSource" -> provider)) From a6f6e07b70311fb843670b89f6546ae675359feb Mon Sep 17 00:00:00 2001 From: Yuchen Liu Date: Wed, 18 Sep 2024 15:45:17 -0700 Subject: [PATCH 057/189] [SPARK-48939][AVRO] Support reading Avro with recursive schema reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Continue the discussion from https://github.com/apache/spark/pull/47425 to this PR because I can't push to Yuchen's account ### What changes were proposed in this pull request? The builtin ProtoBuf connector first supports recursive schema reference. It is approached by letting users specify an option “recursive.fields.max.depth”, and at the start of the execution, unroll the recursive field by this level. It converts a problem of dynamic schema for each row to a fixed schema which is supported by Spark. Avro can just adopt a similar method. This PR defines an option "recursiveFieldMaxDepth" to both Avro data source and from_avro function. With this option, Spark can support Avro recursive schema up to certain depth. ### Why are the changes needed? Recursive reference denotes the case that the type of a field can be defined before in the parent nodes. A simple example is: ``` { "type": "record", "name": "LongList", "fields" : [ {"name": "value", "type": "long"}, {"name": "next", "type": ["null", "LongList"]} ] } ``` This is written in Avro Schema DSL and represents a linked list data structure. Spark currently will throw an error on this schema. Many users used schema like this, so we should support it. ### Does this PR introduce any user-facing change? Yes. Previously, it will throw error on recursive schemas like above. With this change, it will still throw the same error by default but when users specify the option to a number greater than 0, the schema will be unrolled to that depth. ### How was this patch tested? Added new unit tests and integration tests to AvroSuite and AvroFunctionSuite. ### Was this patch authored or co-authored using generative AI tooling? No. Co-authored-by: Wei Liu Closes #48043 from WweiL/yuchen-avro-recursive-schema. Lead-authored-by: Yuchen Liu Co-authored-by: Wei Liu Co-authored-by: Yuchen Liu <170372783+eason-yuchen-liu@users.noreply.github.com> Signed-off-by: Gengliang Wang --- .../org/apache/spark/internal/LogKey.scala | 2 + .../spark/sql/avro/AvroDataToCatalyst.scala | 6 +- .../spark/sql/avro/AvroDeserializer.scala | 12 +- .../spark/sql/avro/AvroFileFormat.scala | 3 +- .../apache/spark/sql/avro/AvroOptions.scala | 31 +++ .../org/apache/spark/sql/avro/AvroUtils.scala | 3 +- .../spark/sql/avro/SchemaConverters.scala | 198 +++++++++++----- .../v2/avro/AvroPartitionReaderFactory.scala | 3 +- .../AvroCatalystDataConversionSuite.scala | 3 +- .../spark/sql/avro/AvroFunctionsSuite.scala | 33 ++- .../spark/sql/avro/AvroRowReaderSuite.scala | 3 +- .../spark/sql/avro/AvroSerdeSuite.scala | 3 +- .../org/apache/spark/sql/avro/AvroSuite.scala | 223 +++++++++++++++++- docs/sql-data-sources-avro.md | 45 ++++ .../sql/errors/QueryCompilationErrors.scala | 7 + 15 files changed, 488 insertions(+), 87 deletions(-) diff --git a/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala b/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala index a7e4f186000b5..12d456a371d07 100644 --- a/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala +++ b/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala @@ -266,6 +266,7 @@ private[spark] object LogKeys { case object FEATURE_NAME extends LogKey case object FETCH_SIZE extends LogKey case object FIELD_NAME extends LogKey + case object FIELD_TYPE extends LogKey case object FILES extends LogKey case object FILE_ABSOLUTE_PATH extends LogKey case object FILE_END_OFFSET extends LogKey @@ -652,6 +653,7 @@ private[spark] object LogKeys { case object RECEIVER_IDS extends LogKey case object RECORDS extends LogKey case object RECOVERY_STATE extends LogKey + case object RECURSIVE_DEPTH extends LogKey case object REDACTED_STATEMENT extends LogKey case object REDUCE_ID extends LogKey case object REGEX extends LogKey diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala index 7d80998d96eb1..0b85b208242cb 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDataToCatalyst.scala @@ -42,7 +42,8 @@ private[sql] case class AvroDataToCatalyst( val dt = SchemaConverters.toSqlType( expectedSchema, avroOptions.useStableIdForUnionType, - avroOptions.stableIdPrefixForUnionType).dataType + avroOptions.stableIdPrefixForUnionType, + avroOptions.recursiveFieldMaxDepth).dataType parseMode match { // With PermissiveMode, the output Catalyst row might contain columns of null values for // corrupt records, even if some of the columns are not nullable in the user-provided schema. @@ -69,7 +70,8 @@ private[sql] case class AvroDataToCatalyst( dataType, avroOptions.datetimeRebaseModeInRead, avroOptions.useStableIdForUnionType, - avroOptions.stableIdPrefixForUnionType) + avroOptions.stableIdPrefixForUnionType, + avroOptions.recursiveFieldMaxDepth) @transient private var decoder: BinaryDecoder = _ diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala index 877c3f89e88c0..ac20614553ca2 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala @@ -51,14 +51,16 @@ private[sql] class AvroDeserializer( datetimeRebaseSpec: RebaseSpec, filters: StructFilters, useStableIdForUnionType: Boolean, - stableIdPrefixForUnionType: String) { + stableIdPrefixForUnionType: String, + recursiveFieldMaxDepth: Int) { def this( rootAvroType: Schema, rootCatalystType: DataType, datetimeRebaseMode: String, useStableIdForUnionType: Boolean, - stableIdPrefixForUnionType: String) = { + stableIdPrefixForUnionType: String, + recursiveFieldMaxDepth: Int) = { this( rootAvroType, rootCatalystType, @@ -66,7 +68,8 @@ private[sql] class AvroDeserializer( RebaseSpec(LegacyBehaviorPolicy.withName(datetimeRebaseMode)), new NoopFilters, useStableIdForUnionType, - stableIdPrefixForUnionType) + stableIdPrefixForUnionType, + recursiveFieldMaxDepth) } private lazy val decimalConversions = new DecimalConversion() @@ -128,7 +131,8 @@ private[sql] class AvroDeserializer( s"schema is incompatible (avroType = $avroType, sqlType = ${catalystType.sql})" val realDataType = SchemaConverters.toSqlType( - avroType, useStableIdForUnionType, stableIdPrefixForUnionType).dataType + avroType, useStableIdForUnionType, stableIdPrefixForUnionType, + recursiveFieldMaxDepth).dataType (avroType.getType, catalystType) match { case (NULL, NullType) => (updater, ordinal, _) => diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala index 372f24b54f5c4..264c3a1f48abe 100755 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala @@ -145,7 +145,8 @@ private[sql] class AvroFileFormat extends FileFormat datetimeRebaseMode, avroFilters, parsedOptions.useStableIdForUnionType, - parsedOptions.stableIdPrefixForUnionType) + parsedOptions.stableIdPrefixForUnionType, + parsedOptions.recursiveFieldMaxDepth) override val stopPosition = file.start + file.length override def hasNext: Boolean = hasNextRow diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala index 4332904339f19..e0c6ad3ee69d3 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroOptions.scala @@ -27,6 +27,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.{DataSourceOptions, FileSourceOptions} import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, FailFastMode, ParseMode} +import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf /** @@ -136,6 +137,15 @@ private[sql] class AvroOptions( val stableIdPrefixForUnionType: String = parameters .getOrElse(STABLE_ID_PREFIX_FOR_UNION_TYPE, "member_") + + val recursiveFieldMaxDepth: Int = + parameters.get(RECURSIVE_FIELD_MAX_DEPTH).map(_.toInt).getOrElse(-1) + + if (recursiveFieldMaxDepth > RECURSIVE_FIELD_MAX_DEPTH_LIMIT) { + throw QueryCompilationErrors.avroOptionsException( + RECURSIVE_FIELD_MAX_DEPTH, + s"Should not be greater than $RECURSIVE_FIELD_MAX_DEPTH_LIMIT.") + } } private[sql] object AvroOptions extends DataSourceOptions { @@ -170,4 +180,25 @@ private[sql] object AvroOptions extends DataSourceOptions { // When STABLE_ID_FOR_UNION_TYPE is enabled, the option allows to configure the prefix for fields // of Avro Union type. val STABLE_ID_PREFIX_FOR_UNION_TYPE = newOption("stableIdentifierPrefixForUnionType") + + /** + * Adds support for recursive fields. If this option is not specified or is set to 0, recursive + * fields are not permitted. Setting it to 1 drops all recursive fields, 2 allows recursive + * fields to be recursed once, and 3 allows it to be recursed twice and so on, up to 15. + * Values larger than 15 are not allowed in order to avoid inadvertently creating very large + * schemas. If an avro message has depth beyond this limit, the Spark struct returned is + * truncated after the recursion limit. + * + * Examples: Consider an Avro schema with a recursive field: + * {"type" : "record", "name" : "Node", "fields" : [{"name": "Id", "type": "int"}, + * {"name": "Next", "type": ["null", "Node"]}]} + * The following lists the parsed schema with different values for this setting. + * 1: `struct` + * 2: `struct>` + * 3: `struct>>` + * and so on. + */ + val RECURSIVE_FIELD_MAX_DEPTH = newOption("recursiveFieldMaxDepth") + + val RECURSIVE_FIELD_MAX_DEPTH_LIMIT: Int = 15 } diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala index 7cbc30f1fb3dc..594ebb4716c41 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala @@ -65,7 +65,8 @@ private[sql] object AvroUtils extends Logging { SchemaConverters.toSqlType( avroSchema, parsedOptions.useStableIdForUnionType, - parsedOptions.stableIdPrefixForUnionType).dataType match { + parsedOptions.stableIdPrefixForUnionType, + parsedOptions.recursiveFieldMaxDepth).dataType match { case t: StructType => Some(t) case _ => throw new RuntimeException( s"""Avro schema cannot be converted to a Spark SQL StructType: diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala index b2285aa966ddb..1168a887abd8e 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala @@ -27,6 +27,10 @@ import org.apache.avro.LogicalTypes.{Date, Decimal, LocalTimestampMicros, LocalT import org.apache.avro.Schema.Type._ import org.apache.spark.annotation.DeveloperApi +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{FIELD_NAME, FIELD_TYPE, RECURSIVE_DEPTH} +import org.apache.spark.internal.MDC +import org.apache.spark.sql.avro.AvroOptions.RECURSIVE_FIELD_MAX_DEPTH_LIMIT import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.types._ import org.apache.spark.sql.types.Decimal.minBytesForPrecision @@ -36,7 +40,7 @@ import org.apache.spark.sql.types.Decimal.minBytesForPrecision * versa. */ @DeveloperApi -object SchemaConverters { +object SchemaConverters extends Logging { private lazy val nullSchema = Schema.create(Schema.Type.NULL) /** @@ -48,14 +52,27 @@ object SchemaConverters { /** * Converts an Avro schema to a corresponding Spark SQL schema. - * + * + * @param avroSchema The Avro schema to convert. + * @param useStableIdForUnionType If true, Avro schema is deserialized into Spark SQL schema, + * and the Avro Union type is transformed into a structure where + * the field names remain consistent with their respective types. + * @param stableIdPrefixForUnionType The prefix to use to configure the prefix for fields of + * Avro Union type + * @param recursiveFieldMaxDepth The maximum depth to recursively process fields in Avro schema. + * -1 means not supported. * @since 4.0.0 */ def toSqlType( avroSchema: Schema, useStableIdForUnionType: Boolean, - stableIdPrefixForUnionType: String): SchemaType = { - toSqlTypeHelper(avroSchema, Set.empty, useStableIdForUnionType, stableIdPrefixForUnionType) + stableIdPrefixForUnionType: String, + recursiveFieldMaxDepth: Int = -1): SchemaType = { + val schema = toSqlTypeHelper(avroSchema, Map.empty, useStableIdForUnionType, + stableIdPrefixForUnionType, recursiveFieldMaxDepth) + // the top level record should never return null + assert(schema != null) + schema } /** * Converts an Avro schema to a corresponding Spark SQL schema. @@ -63,17 +80,17 @@ object SchemaConverters { * @since 2.4.0 */ def toSqlType(avroSchema: Schema): SchemaType = { - toSqlType(avroSchema, false, "") + toSqlType(avroSchema, false, "", -1) } @deprecated("using toSqlType(..., useStableIdForUnionType: Boolean) instead", "4.0.0") def toSqlType(avroSchema: Schema, options: Map[String, String]): SchemaType = { val avroOptions = AvroOptions(options) - toSqlTypeHelper( + toSqlType( avroSchema, - Set.empty, avroOptions.useStableIdForUnionType, - avroOptions.stableIdPrefixForUnionType) + avroOptions.stableIdPrefixForUnionType, + avroOptions.recursiveFieldMaxDepth) } // The property specifies Catalyst type of the given field @@ -81,9 +98,10 @@ object SchemaConverters { private def toSqlTypeHelper( avroSchema: Schema, - existingRecordNames: Set[String], + existingRecordNames: Map[String, Int], useStableIdForUnionType: Boolean, - stableIdPrefixForUnionType: String): SchemaType = { + stableIdPrefixForUnionType: String, + recursiveFieldMaxDepth: Int): SchemaType = { avroSchema.getType match { case INT => avroSchema.getLogicalType match { case _: Date => SchemaType(DateType, nullable = false) @@ -128,62 +146,110 @@ object SchemaConverters { case NULL => SchemaType(NullType, nullable = true) case RECORD => - if (existingRecordNames.contains(avroSchema.getFullName)) { + val recursiveDepth: Int = existingRecordNames.getOrElse(avroSchema.getFullName, 0) + if (recursiveDepth > 0 && recursiveFieldMaxDepth <= 0) { throw new IncompatibleSchemaException(s""" - |Found recursive reference in Avro schema, which can not be processed by Spark: - |${avroSchema.toString(true)} + |Found recursive reference in Avro schema, which can not be processed by Spark by + | default: ${avroSchema.toString(true)}. Try setting the option `recursiveFieldMaxDepth` + | to 1 - $RECURSIVE_FIELD_MAX_DEPTH_LIMIT. """.stripMargin) - } - val newRecordNames = existingRecordNames + avroSchema.getFullName - val fields = avroSchema.getFields.asScala.map { f => - val schemaType = toSqlTypeHelper( - f.schema(), - newRecordNames, - useStableIdForUnionType, - stableIdPrefixForUnionType) - StructField(f.name, schemaType.dataType, schemaType.nullable) - } + } else if (recursiveDepth > 0 && recursiveDepth >= recursiveFieldMaxDepth) { + logInfo( + log"The field ${MDC(FIELD_NAME, avroSchema.getFullName)} of type " + + log"${MDC(FIELD_TYPE, avroSchema.getType.getName)} is dropped at recursive depth " + + log"${MDC(RECURSIVE_DEPTH, recursiveDepth)}." + ) + null + } else { + val newRecordNames = + existingRecordNames + (avroSchema.getFullName -> (recursiveDepth + 1)) + val fields = avroSchema.getFields.asScala.map { f => + val schemaType = toSqlTypeHelper( + f.schema(), + newRecordNames, + useStableIdForUnionType, + stableIdPrefixForUnionType, + recursiveFieldMaxDepth) + if (schemaType == null) { + null + } + else { + StructField(f.name, schemaType.dataType, schemaType.nullable) + } + }.filter(_ != null).toSeq - SchemaType(StructType(fields.toArray), nullable = false) + SchemaType(StructType(fields), nullable = false) + } case ARRAY => val schemaType = toSqlTypeHelper( avroSchema.getElementType, existingRecordNames, useStableIdForUnionType, - stableIdPrefixForUnionType) - SchemaType( - ArrayType(schemaType.dataType, containsNull = schemaType.nullable), - nullable = false) + stableIdPrefixForUnionType, + recursiveFieldMaxDepth) + if (schemaType == null) { + logInfo( + log"Dropping ${MDC(FIELD_NAME, avroSchema.getFullName)} of type " + + log"${MDC(FIELD_TYPE, avroSchema.getType.getName)} as it does not have any " + + log"fields left likely due to recursive depth limit." + ) + null + } else { + SchemaType( + ArrayType(schemaType.dataType, containsNull = schemaType.nullable), + nullable = false) + } case MAP => val schemaType = toSqlTypeHelper(avroSchema.getValueType, - existingRecordNames, useStableIdForUnionType, stableIdPrefixForUnionType) - SchemaType( - MapType(StringType, schemaType.dataType, valueContainsNull = schemaType.nullable), - nullable = false) + existingRecordNames, useStableIdForUnionType, stableIdPrefixForUnionType, + recursiveFieldMaxDepth) + if (schemaType == null) { + logInfo( + log"Dropping ${MDC(FIELD_NAME, avroSchema.getFullName)} of type " + + log"${MDC(FIELD_TYPE, avroSchema.getType.getName)} as it does not have any " + + log"fields left likely due to recursive depth limit." + ) + null + } else { + SchemaType( + MapType(StringType, schemaType.dataType, valueContainsNull = schemaType.nullable), + nullable = false) + } case UNION => if (avroSchema.getTypes.asScala.exists(_.getType == NULL)) { // In case of a union with null, eliminate it and make a recursive call val remainingUnionTypes = AvroUtils.nonNullUnionBranches(avroSchema) - if (remainingUnionTypes.size == 1) { - toSqlTypeHelper( - remainingUnionTypes.head, - existingRecordNames, - useStableIdForUnionType, - stableIdPrefixForUnionType).copy(nullable = true) + val remainingSchema = + if (remainingUnionTypes.size == 1) { + remainingUnionTypes.head + } else { + Schema.createUnion(remainingUnionTypes.asJava) + } + val schemaType = toSqlTypeHelper( + remainingSchema, + existingRecordNames, + useStableIdForUnionType, + stableIdPrefixForUnionType, + recursiveFieldMaxDepth) + + if (schemaType == null) { + logInfo( + log"Dropping ${MDC(FIELD_NAME, avroSchema.getFullName)} of type " + + log"${MDC(FIELD_TYPE, avroSchema.getType.getName)} as it does not have any " + + log"fields left likely due to recursive depth limit." + ) + null } else { - toSqlTypeHelper( - Schema.createUnion(remainingUnionTypes.asJava), - existingRecordNames, - useStableIdForUnionType, - stableIdPrefixForUnionType).copy(nullable = true) + schemaType.copy(nullable = true) } } else avroSchema.getTypes.asScala.map(_.getType).toSeq match { case Seq(t1) => toSqlTypeHelper(avroSchema.getTypes.get(0), - existingRecordNames, useStableIdForUnionType, stableIdPrefixForUnionType) + existingRecordNames, useStableIdForUnionType, stableIdPrefixForUnionType, + recursiveFieldMaxDepth) case Seq(t1, t2) if Set(t1, t2) == Set(INT, LONG) => SchemaType(LongType, nullable = false) case Seq(t1, t2) if Set(t1, t2) == Set(FLOAT, DOUBLE) => @@ -201,29 +267,33 @@ object SchemaConverters { s, existingRecordNames, useStableIdForUnionType, - stableIdPrefixForUnionType) - - val fieldName = if (useStableIdForUnionType) { - // Avro's field name may be case sensitive, so field names for two named type - // could be "a" and "A" and we need to distinguish them. In this case, we throw - // an exception. - // Stable id prefix can be empty so the name of the field can be just the type. - val tempFieldName = s"${stableIdPrefixForUnionType}${s.getName}" - if (!fieldNameSet.add(tempFieldName.toLowerCase(Locale.ROOT))) { - throw new IncompatibleSchemaException( - "Cannot generate stable identifier for Avro union type due to name " + - s"conflict of type name ${s.getName}") - } - tempFieldName + stableIdPrefixForUnionType, + recursiveFieldMaxDepth) + if (schemaType == null) { + null } else { - s"member$i" - } + val fieldName = if (useStableIdForUnionType) { + // Avro's field name may be case sensitive, so field names for two named type + // could be "a" and "A" and we need to distinguish them. In this case, we throw + // an exception. + // Stable id prefix can be empty so the name of the field can be just the type. + val tempFieldName = s"${stableIdPrefixForUnionType}${s.getName}" + if (!fieldNameSet.add(tempFieldName.toLowerCase(Locale.ROOT))) { + throw new IncompatibleSchemaException( + "Cannot generate stable identifier for Avro union type due to name " + + s"conflict of type name ${s.getName}") + } + tempFieldName + } else { + s"member$i" + } - // All fields are nullable because only one of them is set at a time - StructField(fieldName, schemaType.dataType, nullable = true) - } + // All fields are nullable because only one of them is set at a time + StructField(fieldName, schemaType.dataType, nullable = true) + } + }.filter(_ != null).toSeq - SchemaType(StructType(fields.toArray), nullable = false) + SchemaType(StructType(fields), nullable = false) } case other => throw new IncompatibleSchemaException(s"Unsupported type $other") diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroPartitionReaderFactory.scala b/connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroPartitionReaderFactory.scala index 1083c99160724..a13faf3b51560 100644 --- a/connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroPartitionReaderFactory.scala +++ b/connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroPartitionReaderFactory.scala @@ -105,7 +105,8 @@ case class AvroPartitionReaderFactory( datetimeRebaseMode, avroFilters, options.useStableIdForUnionType, - options.stableIdPrefixForUnionType) + options.stableIdPrefixForUnionType, + options.recursiveFieldMaxDepth) override val stopPosition = partitionedFile.start + partitionedFile.length override def next(): Boolean = hasNextRow diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala index 388347537a4d6..311eda3a1b6ae 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala @@ -291,7 +291,8 @@ class AvroCatalystDataConversionSuite extends SparkFunSuite RebaseSpec(LegacyBehaviorPolicy.CORRECTED), filters, false, - "") + "", + -1) val deserialized = deserializer.deserialize(data) expected match { case None => assert(deserialized == None) diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroFunctionsSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroFunctionsSuite.scala index 47faaf7662a50..a7f7abadcf485 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroFunctionsSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroFunctionsSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.LocalTableScanExec import org.apache.spark.sql.functions.{col, lit, struct} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.sql.types.{BinaryType, StructType} +import org.apache.spark.sql.types.{BinaryType, IntegerType, StructField, StructType} class AvroFunctionsSuite extends QueryTest with SharedSparkSession { import testImplicits._ @@ -374,6 +374,37 @@ class AvroFunctionsSuite extends QueryTest with SharedSparkSession { } } + + test("roundtrip in to_avro and from_avro - recursive schema") { + val catalystSchema = + StructType(Seq( + StructField("Id", IntegerType), + StructField("Name", StructType(Seq( + StructField("Id", IntegerType), + StructField("Name", StructType(Seq( + StructField("Id", IntegerType))))))))) + + val avroSchema = s""" + |{ + | "type" : "record", + | "name" : "test_schema", + | "fields" : [ + | {"name": "Id", "type": "int"}, + | {"name": "Name", "type": ["null", "test_schema"]} + | ] + |} + """.stripMargin + + val df = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(2, Row(3, Row(4))), Row(1, null))), + catalystSchema).select(struct("Id", "Name").as("struct")) + + val avroStructDF = df.select(functions.to_avro($"struct", avroSchema).as("avro")) + checkAnswer(avroStructDF.select( + functions.from_avro($"avro", avroSchema, Map( + "recursiveFieldMaxDepth" -> "3").asJava)), df) + } + private def serialize(record: GenericRecord, avroSchema: String): Array[Byte] = { val schema = new Schema.Parser().parse(avroSchema) val datumWriter = new GenericDatumWriter[GenericRecord](schema) diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroRowReaderSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroRowReaderSuite.scala index 9b3bb929a700d..c1ab96a63eb26 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroRowReaderSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroRowReaderSuite.scala @@ -77,7 +77,8 @@ class AvroRowReaderSuite RebaseSpec(CORRECTED), new NoopFilters, false, - "") + "", + -1) override val stopPosition = fileSize override def hasNext: Boolean = hasNextRow diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSerdeSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSerdeSuite.scala index cbcbc2e7e76a6..3643a95abe19c 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSerdeSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSerdeSuite.scala @@ -228,7 +228,8 @@ object AvroSerdeSuite { RebaseSpec(CORRECTED), new NoopFilters, false, - "") + "", + -1) } /** diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSuite.scala index 14ed6c43e4c0f..be887bd5237b0 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroSuite.scala @@ -2220,7 +2220,8 @@ abstract class AvroSuite } } - private def checkSchemaWithRecursiveLoop(avroSchema: String): Unit = { + private def checkSchemaWithRecursiveLoop(avroSchema: String, recursiveFieldMaxDepth: Int): + Unit = { val message = intercept[IncompatibleSchemaException] { SchemaConverters.toSqlType(new Schema.Parser().parse(avroSchema), false, "") }.getMessage @@ -2229,7 +2230,79 @@ abstract class AvroSuite } test("Detect recursive loop") { - checkSchemaWithRecursiveLoop(""" + for (recursiveFieldMaxDepth <- Seq(-1, 0)) { + checkSchemaWithRecursiveLoop( + """ + |{ + | "type": "record", + | "name": "LongList", + | "fields" : [ + | {"name": "value", "type": "long"}, // each element has a long + | {"name": "next", "type": ["null", "LongList"]} // optional next element + | ] + |} + """.stripMargin, recursiveFieldMaxDepth) + + checkSchemaWithRecursiveLoop( + """ + |{ + | "type": "record", + | "name": "LongList", + | "fields": [ + | { + | "name": "value", + | "type": { + | "type": "record", + | "name": "foo", + | "fields": [ + | { + | "name": "parent", + | "type": "LongList" + | } + | ] + | } + | } + | ] + |} + """.stripMargin, recursiveFieldMaxDepth) + + checkSchemaWithRecursiveLoop( + """ + |{ + | "type": "record", + | "name": "LongList", + | "fields" : [ + | {"name": "value", "type": "long"}, + | {"name": "array", "type": {"type": "array", "items": "LongList"}} + | ] + |} + """.stripMargin, recursiveFieldMaxDepth) + + checkSchemaWithRecursiveLoop( + """ + |{ + | "type": "record", + | "name": "LongList", + | "fields" : [ + | {"name": "value", "type": "long"}, + | {"name": "map", "type": {"type": "map", "values": "LongList"}} + | ] + |} + """.stripMargin, recursiveFieldMaxDepth) + } + } + + private def checkSparkSchemaEquals( + avroSchema: String, expectedSchema: StructType, recursiveFieldMaxDepth: Int): Unit = { + val sparkSchema = + SchemaConverters.toSqlType( + new Schema.Parser().parse(avroSchema), false, "", recursiveFieldMaxDepth).dataType + + assert(sparkSchema === expectedSchema) + } + + test("Translate recursive schema - union") { + val avroSchema = """ |{ | "type": "record", | "name": "LongList", @@ -2238,9 +2311,57 @@ abstract class AvroSuite | {"name": "next", "type": ["null", "LongList"]} // optional next element | ] |} - """.stripMargin) + """.stripMargin + val nonRecursiveFields = new StructType().add("value", LongType, nullable = false) + var expectedSchema = nonRecursiveFields + for (i <- 1 to 5) { + checkSparkSchemaEquals(avroSchema, expectedSchema, i) + expectedSchema = nonRecursiveFields.add("next", expectedSchema) + } + } + + test("Translate recursive schema - union - 2 non-null fields") { + val avroSchema = """ + |{ + | "type": "record", + | "name": "TreeNode", + | "fields": [ + | { + | "name": "name", + | "type": "string" + | }, + | { + | "name": "value", + | "type": [ + | "long" + | ] + | }, + | { + | "name": "children", + | "type": [ + | "null", + | { + | "type": "array", + | "items": "TreeNode" + | } + | ], + | "default": null + | } + | ] + |} + """.stripMargin + val nonRecursiveFields = new StructType().add("name", StringType, nullable = false) + .add("value", LongType, nullable = false) + var expectedSchema = nonRecursiveFields + for (i <- 1 to 5) { + checkSparkSchemaEquals(avroSchema, expectedSchema, i) + expectedSchema = nonRecursiveFields.add("children", + new ArrayType(expectedSchema, false), nullable = true) + } + } - checkSchemaWithRecursiveLoop(""" + test("Translate recursive schema - record") { + val avroSchema = """ |{ | "type": "record", | "name": "LongList", @@ -2260,9 +2381,18 @@ abstract class AvroSuite | } | ] |} - """.stripMargin) + """.stripMargin + val nonRecursiveFields = new StructType().add("value", StructType(Seq()), nullable = false) + var expectedSchema = nonRecursiveFields + for (i <- 1 to 5) { + checkSparkSchemaEquals(avroSchema, expectedSchema, i) + expectedSchema = new StructType().add("value", + new StructType().add("parent", expectedSchema, nullable = false), nullable = false) + } + } - checkSchemaWithRecursiveLoop(""" + test("Translate recursive schema - array") { + val avroSchema = """ |{ | "type": "record", | "name": "LongList", @@ -2271,9 +2401,18 @@ abstract class AvroSuite | {"name": "array", "type": {"type": "array", "items": "LongList"}} | ] |} - """.stripMargin) + """.stripMargin + val nonRecursiveFields = new StructType().add("value", LongType, nullable = false) + var expectedSchema = nonRecursiveFields + for (i <- 1 to 5) { + checkSparkSchemaEquals(avroSchema, expectedSchema, i) + expectedSchema = + nonRecursiveFields.add("array", new ArrayType(expectedSchema, false), nullable = false) + } + } - checkSchemaWithRecursiveLoop(""" + test("Translate recursive schema - map") { + val avroSchema = """ |{ | "type": "record", | "name": "LongList", @@ -2282,7 +2421,70 @@ abstract class AvroSuite | {"name": "map", "type": {"type": "map", "values": "LongList"}} | ] |} - """.stripMargin) + """.stripMargin + val nonRecursiveFields = new StructType().add("value", LongType, nullable = false) + var expectedSchema = nonRecursiveFields + for (i <- 1 to 5) { + checkSparkSchemaEquals(avroSchema, expectedSchema, i) + expectedSchema = + nonRecursiveFields.add("map", + new MapType(StringType, expectedSchema, false), nullable = false) + } + } + + test("recursive schema integration test") { + val catalystSchema = + StructType(Seq( + StructField("Id", IntegerType), + StructField("Name", StructType(Seq( + StructField("Id", IntegerType), + StructField("Name", StructType(Seq( + StructField("Id", IntegerType), + StructField("Name", NullType))))))))) + + val avroSchema = s""" + |{ + | "type" : "record", + | "name" : "test_schema", + | "fields" : [ + | {"name": "Id", "type": "int"}, + | {"name": "Name", "type": ["null", "test_schema"]} + | ] + |} + """.stripMargin + + val df = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(2, Row(3, Row(4, null))), Row(1, null))), + catalystSchema) + + withTempPath { tempDir => + df.write.format("avro").save(tempDir.getPath) + + val exc = intercept[AnalysisException] { + spark.read + .format("avro") + .option("avroSchema", avroSchema) + .option("recursiveFieldMaxDepth", 16) + .load(tempDir.getPath) + } + assert(exc.getMessage.contains("Should not be greater than 15.")) + + checkAnswer( + spark.read + .format("avro") + .option("avroSchema", avroSchema) + .option("recursiveFieldMaxDepth", 10) + .load(tempDir.getPath), + df) + + checkAnswer( + spark.read + .format("avro") + .option("avroSchema", avroSchema) + .option("recursiveFieldMaxDepth", 1) + .load(tempDir.getPath), + df.select("Id")) + } } test("log a warning of ignoreExtension deprecation") { @@ -2777,7 +2979,7 @@ abstract class AvroSuite } test("SPARK-40667: validate Avro Options") { - assert(AvroOptions.getAllOptions.size == 11) + assert(AvroOptions.getAllOptions.size == 12) // Please add validation on any new Avro options here assert(AvroOptions.isValidOption("ignoreExtension")) assert(AvroOptions.isValidOption("mode")) @@ -2790,6 +2992,7 @@ abstract class AvroSuite assert(AvroOptions.isValidOption("datetimeRebaseMode")) assert(AvroOptions.isValidOption("enableStableIdentifiersForUnionType")) assert(AvroOptions.isValidOption("stableIdentifierPrefixForUnionType")) + assert(AvroOptions.isValidOption("recursiveFieldMaxDepth")) } test("SPARK-46633: read file with empty blocks") { diff --git a/docs/sql-data-sources-avro.md b/docs/sql-data-sources-avro.md index 3721f92d93266..c06e1fd46d2da 100644 --- a/docs/sql-data-sources-avro.md +++ b/docs/sql-data-sources-avro.md @@ -353,6 +353,13 @@ Data source options of Avro can be set via: read 4.0.0 + + recursiveFieldMaxDepth + -1 + If this option is specified to negative or is set to 0, recursive fields are not permitted. Setting it to 1 drops all recursive fields, 2 allows recursive fields to be recursed once, and 3 allows it to be recursed twice and so on, up to 15. Values larger than 15 are not allowed in order to avoid inadvertently creating very large schemas. If an avro message has depth beyond this limit, the Spark struct returned is truncated after the recursion limit. An example of usage can be found in section Handling circular references of Avro fields + read + 4.0.0 + ## Configuration @@ -628,3 +635,41 @@ You can also specify the whole output Avro schema with the option `avroSchema`, decimal + +## Handling circular references of Avro fields +In Avro, a circular reference occurs when the type of a field is defined in one of the parent records. This can cause issues when parsing the data, as it can result in infinite loops or other unexpected behavior. +To read Avro data with schema that has circular reference, users can use the `recursiveFieldMaxDepth` option to specify the maximum number of levels of recursion to allow when parsing the schema. By default, Spark Avro data source will not permit recursive fields by setting `recursiveFieldMaxDepth` to -1. However, you can set this option to 1 to 15 if needed. + +Setting `recursiveFieldMaxDepth` to 1 drops all recursive fields, setting it to 2 allows it to be recursed once, and setting it to 3 allows it to be recursed twice. A `recursiveFieldMaxDepth` value greater than 15 is not allowed, as it can lead to performance issues and even stack overflows. + +SQL Schema for the below Avro message will vary based on the value of `recursiveFieldMaxDepth`. + +

+
+This div is only used to make markdown editor/viewer happy and does not display on web + +```avro +
+ +{% highlight avro %} +{ + "type": "record", + "name": "Node", + "fields": [ + {"name": "Id", "type": "int"}, + {"name": "Next", "type": ["null", "Node"]} + ] +} + +// The Avro schema defined above, would be converted into a Spark SQL columns with the following +// structure based on `recursiveFieldMaxDepth` value. + +1: struct +2: struct> +3: struct>> + +{% endhighlight %} +
+``` +
+
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index e324d4e9d2edb..ad0e1d07bf93d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -4090,6 +4090,13 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat ) } + def avroOptionsException(optionName: String, message: String): Throwable = { + new AnalysisException( + errorClass = "STDS_INVALID_OPTION_VALUE.WITH_MESSAGE", + messageParameters = Map("optionName" -> optionName, "message" -> message) + ) + } + def protobufNotLoadedSqlFunctionsUnusable(functionName: String): Throwable = { new AnalysisException( errorClass = "PROTOBUF_NOT_LOADED_SQL_FUNCTIONS_UNUSABLE", From 25d6b7a280f690c1a467f65143115cce846a732a Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 07:46:18 +0800 Subject: [PATCH 058/189] [SPARK-49692][PYTHON][CONNECT] Refine the string representation of literal date and datetime ### What changes were proposed in this pull request? Refine the string representation of literal date and datetime ### Why are the changes needed? 1, we should not represent those literals with internal values; 2, the string representation should be consistent with PySpark Classic if possible (we cannot make sure the representations are always the same because we only hold an unresolved expression in connect, but we can try our best to do so) ### Does this PR introduce _any_ user-facing change? yes before: ``` In [3]: lit(datetime.date(2024, 7, 10)) Out[3]: Column<'19914'> In [4]: lit(datetime.datetime(2024, 7, 10, 1, 2, 3, 456)) Out[4]: Column<'1720544523000456'> ``` after: ``` In [3]: lit(datetime.date(2024, 7, 10)) Out[3]: Column<'2024-07-10'> In [4]: lit(datetime.datetime(2024, 7, 10, 1, 2, 3, 456)) Out[4]: Column<'2024-07-10 01:02:03.000456'> ``` ### How was this patch tested? added tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48137 from zhengruifeng/py_connect_lit_dt. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- python/pyspark/sql/connect/expressions.py | 16 ++++++++++++++-- python/pyspark/sql/tests/test_column.py | 9 +++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/connect/expressions.py b/python/pyspark/sql/connect/expressions.py index db1cd1c013be5..63128ef48e389 100644 --- a/python/pyspark/sql/connect/expressions.py +++ b/python/pyspark/sql/connect/expressions.py @@ -477,8 +477,20 @@ def to_plan(self, session: "SparkConnectClient") -> "proto.Expression": def __repr__(self) -> str: if self._value is None: return "NULL" - else: - return f"{self._value}" + elif isinstance(self._dataType, DateType): + dt = DateType().fromInternal(self._value) + if dt is not None and isinstance(dt, datetime.date): + return dt.strftime("%Y-%m-%d") + elif isinstance(self._dataType, TimestampType): + ts = TimestampType().fromInternal(self._value) + if ts is not None and isinstance(ts, datetime.datetime): + return ts.strftime("%Y-%m-%d %H:%M:%S.%f") + elif isinstance(self._dataType, TimestampNTZType): + ts = TimestampNTZType().fromInternal(self._value) + if ts is not None and isinstance(ts, datetime.datetime): + return ts.strftime("%Y-%m-%d %H:%M:%S.%f") + # TODO(SPARK-49693): Refine the string representation of timedelta + return f"{self._value}" class ColumnReference(Expression): diff --git a/python/pyspark/sql/tests/test_column.py b/python/pyspark/sql/tests/test_column.py index 2bd66baaa2bfe..220ecd387f7ee 100644 --- a/python/pyspark/sql/tests/test_column.py +++ b/python/pyspark/sql/tests/test_column.py @@ -18,6 +18,8 @@ from enum import Enum from itertools import chain +import datetime + from pyspark.sql import Column, Row from pyspark.sql import functions as sf from pyspark.sql.types import StructType, StructField, IntegerType, LongType @@ -280,6 +282,13 @@ def test_expr_str_representation(self): when_cond = sf.when(expression, sf.lit(None)) self.assertEqual(str(when_cond), "Column<'CASE WHEN foo THEN NULL END'>") + def test_lit_time_representation(self): + dt = datetime.date(2021, 3, 4) + self.assertEqual(str(sf.lit(dt)), "Column<'2021-03-04'>") + + ts = datetime.datetime(2021, 3, 4, 12, 34, 56, 1234) + self.assertEqual(str(sf.lit(ts)), "Column<'2021-03-04 12:34:56.001234'>") + def test_enum_literals(self): class IntEnum(Enum): X = 1 From 669e63a34012404d8d864cd6294f799b672f6f9a Mon Sep 17 00:00:00 2001 From: Robert Dillitz Date: Thu, 19 Sep 2024 08:54:20 +0900 Subject: [PATCH 059/189] [SPARK-49673][CONNECT] Increase CONNECT_GRPC_ARROW_MAX_BATCH_SIZE to 0.7 * CONNECT_GRPC_MAX_MESSAGE_SIZE ### What changes were proposed in this pull request? Increases the default `maxBatchSize` from 4MiB * 0.7 to 128MiB (= CONNECT_GRPC_MAX_MESSAGE_SIZE) * 0.7. This makes better use of the allowed maximum message size. This limit is used when creating Arrow batches for the `SqlCommandResult` in the `SparkConnectPlanner` and for `ExecutePlanResponse.ArrowBatch` in `processAsArrowBatches`. This, for example, lets us return much larger `LocalRelations` in the `SqlCommandResult` (i.e., for the `SHOW PARTITIONS` command) while still staying within the GRPC message size limit. ### Why are the changes needed? There are `SqlCommandResults` that exceed 0.7 * 4MiB. ### Does this PR introduce _any_ user-facing change? Now support `SqlCommandResults` <= 0.7 * 128 MiB instead of only <= 0.7 * 4MiB and ExecutePlanResponses will now better use the limit of 128MiB. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48122 from dillitz/increase-sql-command-batch-size. Authored-by: Robert Dillitz Signed-off-by: Hyukjin Kwon --- .../apache/spark/sql/ClientE2ETestSuite.scala | 23 +++++++++++++++++-- .../spark/sql/test/RemoteSparkSession.scala | 2 ++ .../spark/sql/connect/config/Connect.scala | 2 +- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala index 52cdbd47357f3..b47231948dc98 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala @@ -23,7 +23,7 @@ import java.util.Properties import scala.collection.mutable import scala.concurrent.{ExecutionContext, Future} -import scala.concurrent.duration.DurationInt +import scala.concurrent.duration.{DurationInt, FiniteDuration} import scala.jdk.CollectionConverters._ import org.apache.commons.io.FileUtils @@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.{NamespaceAlreadyExistsException, import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.StringEncoder import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema import org.apache.spark.sql.catalyst.parser.ParseException -import org.apache.spark.sql.connect.client.{SparkConnectClient, SparkResult} +import org.apache.spark.sql.connect.client.{RetryPolicy, SparkConnectClient, SparkResult} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SqlApiConf import org.apache.spark.sql.test.{ConnectFunSuite, IntegrationTestUtils, RemoteSparkSession, SQLHelper} @@ -1566,6 +1566,25 @@ class ClientE2ETestSuite val result = df.select(trim(col("col"), " ").as("trimmed_col")).collect() assert(result sameElements Array(Row("a"), Row("b"), Row("c"))) } + + test("SPARK-49673: new batch size, multiple batches") { + val maxBatchSize = spark.conf.get("spark.connect.grpc.arrow.maxBatchSize").dropRight(1).toInt + // Adjust client grpcMaxMessageSize to maxBatchSize (10MiB; set in RemoteSparkSession config) + val sparkWithLowerMaxMessageSize = SparkSession + .builder() + .client( + SparkConnectClient + .builder() + .userId("test") + .port(port) + .grpcMaxMessageSize(maxBatchSize) + .retryPolicy(RetryPolicy + .defaultPolicy() + .copy(maxRetries = Some(10), maxBackoff = Some(FiniteDuration(30, "s")))) + .build()) + .create() + assert(sparkWithLowerMaxMessageSize.range(maxBatchSize).collect().length == maxBatchSize) + } } private[sql] case class ClassData(a: String, b: Int) diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/test/RemoteSparkSession.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/test/RemoteSparkSession.scala index e0de73e496d95..36aaa2cc7fbf6 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/test/RemoteSparkSession.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/test/RemoteSparkSession.scala @@ -124,6 +124,8 @@ object SparkConnectServerUtils { // to make the tests exercise reattach. "spark.connect.execute.reattachable.senderMaxStreamDuration=1s", "spark.connect.execute.reattachable.senderMaxStreamSize=123", + // Testing SPARK-49673, setting maxBatchSize to 10MiB + s"spark.connect.grpc.arrow.maxBatchSize=${10 * 1024 * 1024}", // Disable UI "spark.ui.enabled=false") Seq("--jars", catalystTestJar) ++ confs.flatMap(v => "--conf" :: v :: Nil) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala index 92709ff29a1ca..b64637f7d2472 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala @@ -63,7 +63,7 @@ object Connect { "conservatively use 70% of it because the size is not accurate but estimated.") .version("3.4.0") .bytesConf(ByteUnit.BYTE) - .createWithDefault(4 * 1024 * 1024) + .createWithDefault(ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE) val CONNECT_GRPC_MAX_INBOUND_MESSAGE_SIZE = buildStaticConf("spark.connect.grpc.maxInboundMessageSize") From 5c48806a2941070e23a81b4e7e4f3225fe341535 Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Thu, 19 Sep 2024 09:08:59 +0900 Subject: [PATCH 060/189] [SPARK-49688][CONNECT][TESTS] Fix a sporadic `SparkConnectServiceSuite` failure ### What changes were proposed in this pull request? Add a short wait loop to ensure that the test pre-condition is met. To be specific, VerifyEvents.executeHolder is set asynchronously by MockSparkListener.onOtherEvent whereas the test assumes that VerifyEvents.executeHolder is always available. ### Why are the changes needed? For smoother development experience. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? SparkConnectServiceSuite. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48142 from changgyoopark-db/SPARK-49688. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../planner/SparkConnectServiceSuite.scala | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala index 579fdb47aef3c..62146f19328a8 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala @@ -871,10 +871,16 @@ class SparkConnectServiceSuite class VerifyEvents(val sparkContext: SparkContext) { val listener: MockSparkListener = new MockSparkListener() val listenerBus = sparkContext.listenerBus + val EVENT_WAIT_TIMEOUT = timeout(10.seconds) val LISTENER_BUS_TIMEOUT = 30000 def executeHolder: ExecuteHolder = { - assert(listener.executeHolder.isDefined) - listener.executeHolder.get + // An ExecuteHolder shall be set eventually through MockSparkListener + Eventually.eventually(EVENT_WAIT_TIMEOUT) { + assert( + listener.executeHolder.isDefined, + s"No events have been posted in $EVENT_WAIT_TIMEOUT") + listener.executeHolder.get + } } def onNext(v: proto.ExecutePlanResponse): Unit = { if (v.hasSchema) { @@ -891,8 +897,10 @@ class SparkConnectServiceSuite def onCompleted(producedRowCount: Option[Long] = None): Unit = { assert(executeHolder.eventsManager.getProducedRowCount == producedRowCount) // The eventsManager is closed asynchronously - Eventually.eventually(timeout(1.seconds)) { - assert(executeHolder.eventsManager.status == ExecuteStatus.Closed) + Eventually.eventually(EVENT_WAIT_TIMEOUT) { + assert( + executeHolder.eventsManager.status == ExecuteStatus.Closed, + s"Execution has not been completed in $EVENT_WAIT_TIMEOUT") } } def onCanceled(): Unit = { From db8010b4c8be6f1c50f35cbde3efa44cd5d45adf Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Wed, 18 Sep 2024 20:10:18 -0400 Subject: [PATCH 061/189] [SPARK-49568][CONNECT][SQL] Remove self type from Dataset ### What changes were proposed in this pull request? This PR removes the self type parameter from Dataset. This turned out to be a bit noisy. The self type is replaced by a combination of covariant return types and abstract types. Abstract types are used when a method takes a Dataset (or a KeyValueGroupedDataset) as an argument. ### Why are the changes needed? The self type made using the classes in sql/api a bit noisy. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48146 from hvanhovell/SPARK-49568. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../spark/sql/DataFrameNaFunctions.scala | 3 +- .../apache/spark/sql/DataFrameReader.scala | 5 +- .../spark/sql/DataFrameStatFunctions.scala | 3 +- .../scala/org/apache/spark/sql/Dataset.scala | 5 +- .../spark/sql/KeyValueGroupedDataset.scala | 4 +- .../spark/sql/RelationalGroupedDataset.scala | 4 +- .../org/apache/spark/sql/SparkSession.scala | 2 +- .../apache/spark/sql/catalog/Catalog.scala | 3 +- .../sql/connect/ConnectConversions.scala | 51 +++ .../spark/sql/streaming/StreamingQuery.scala | 4 +- .../CheckConnectJvmClientCompatibility.scala | 1 + project/MimaExcludes.scala | 2 + project/SparkBuild.scala | 1 + .../org/apache/spark/sql/api/Catalog.scala | 58 ++-- .../spark/sql/api/DataFrameNaFunctions.scala | 65 ++-- .../spark/sql/api/DataFrameReader.scala | 51 +-- .../sql/api/DataFrameStatFunctions.scala | 22 +- .../org/apache/spark/sql/api/Dataset.scala | 299 +++++++++--------- .../sql/api/KeyValueGroupedDataset.scala | 109 +++---- .../sql/api/RelationalGroupedDataset.scala | 44 ++- .../apache/spark/sql/api/SparkSession.scala | 40 +-- .../apache/spark/sql/api/StreamingQuery.scala | 4 +- .../org/apache/spark/sql/functions.scala | 2 +- .../spark/sql/DataFrameNaFunctions.scala | 3 +- .../apache/spark/sql/DataFrameReader.scala | 5 +- .../spark/sql/DataFrameStatFunctions.scala | 3 +- .../scala/org/apache/spark/sql/Dataset.scala | 4 +- .../spark/sql/KeyValueGroupedDataset.scala | 3 +- .../spark/sql/RelationalGroupedDataset.scala | 5 +- .../org/apache/spark/sql/SparkSession.scala | 2 +- .../apache/spark/sql/catalog/Catalog.scala | 3 +- .../sql/classic/ClassicConversions.scala | 50 +++ .../spark/sql/streaming/StreamingQuery.scala | 4 +- 33 files changed, 500 insertions(+), 364 deletions(-) create mode 100644 connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala index c06cbbc0cdb42..3777f82594aae 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala @@ -22,6 +22,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.connect.proto.{NAReplace, Relation} import org.apache.spark.connect.proto.Expression.{Literal => GLiteral} import org.apache.spark.connect.proto.NAReplace.Replacement +import org.apache.spark.sql.connect.ConnectConversions._ /** * Functionality for working with missing data in `DataFrame`s. @@ -29,7 +30,7 @@ import org.apache.spark.connect.proto.NAReplace.Replacement * @since 3.4.0 */ final class DataFrameNaFunctions private[sql] (sparkSession: SparkSession, root: Relation) - extends api.DataFrameNaFunctions[Dataset] { + extends api.DataFrameNaFunctions { import sparkSession.RichColumn override protected def drop(minNonNulls: Option[Int]): Dataset[Row] = diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameReader.scala index c3ee7030424eb..60bacd4e18ede 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameReader.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameReader.scala @@ -23,6 +23,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Stable import org.apache.spark.connect.proto.Parse.ParseFormat +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.connect.common.DataTypeProtoConverter import org.apache.spark.sql.types.StructType @@ -33,8 +34,8 @@ import org.apache.spark.sql.types.StructType * @since 3.4.0 */ @Stable -class DataFrameReader private[sql] (sparkSession: SparkSession) - extends api.DataFrameReader[Dataset] { +class DataFrameReader private[sql] (sparkSession: SparkSession) extends api.DataFrameReader { + type DS[U] = Dataset[U] /** @inheritdoc */ override def format(source: String): this.type = super.format(source) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala index 9f5ada0d7ec35..bb7cfa75a9ab9 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala @@ -22,6 +22,7 @@ import java.{lang => jl, util => ju} import org.apache.spark.connect.proto.{Relation, StatSampleBy} import org.apache.spark.sql.DataFrameStatFunctions.approxQuantileResultEncoder import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{ArrayEncoder, PrimitiveDoubleEncoder} +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.functions.lit /** @@ -30,7 +31,7 @@ import org.apache.spark.sql.functions.lit * @since 3.4.0 */ final class DataFrameStatFunctions private[sql] (protected val df: DataFrame) - extends api.DataFrameStatFunctions[Dataset] { + extends api.DataFrameStatFunctions { private def root: Relation = df.plan.getRoot private val sparkSession: SparkSession = df.sparkSession diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index 519193ebd9c74..161a0d9d265f0 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.ScalaReflection import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders._ import org.apache.spark.sql.catalyst.expressions.OrderUtils +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.connect.client.SparkResult import org.apache.spark.sql.connect.common.{DataTypeProtoConverter, StorageLevelProtoConverter} import org.apache.spark.sql.errors.DataTypeErrors.toSQLId @@ -134,8 +135,8 @@ class Dataset[T] private[sql] ( val sparkSession: SparkSession, @DeveloperApi val plan: proto.Plan, val encoder: Encoder[T]) - extends api.Dataset[T, Dataset] { - type RGD = RelationalGroupedDataset + extends api.Dataset[T] { + type DS[U] = Dataset[U] import sparkSession.RichColumn diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala index aef7efb08a254..6bf2518901470 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala @@ -26,6 +26,7 @@ import org.apache.spark.api.java.function._ import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.ProductEncoder +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.connect.common.UdfUtils import org.apache.spark.sql.expressions.SparkUserDefinedFunction import org.apache.spark.sql.functions.col @@ -40,8 +41,7 @@ import org.apache.spark.sql.streaming.{GroupState, GroupStateTimeout, OutputMode * * @since 3.5.0 */ -class KeyValueGroupedDataset[K, V] private[sql] () - extends api.KeyValueGroupedDataset[K, V, Dataset] { +class KeyValueGroupedDataset[K, V] private[sql] () extends api.KeyValueGroupedDataset[K, V] { type KVDS[KY, VL] = KeyValueGroupedDataset[KY, VL] private def unsupported(): Nothing = throw new UnsupportedOperationException() diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala index ea13635fc2eaa..14ceb3f4bb144 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql import scala.jdk.CollectionConverters._ import org.apache.spark.connect.proto +import org.apache.spark.sql.connect.ConnectConversions._ /** * A set of methods for aggregations on a `DataFrame`, created by [[Dataset#groupBy groupBy]], @@ -39,8 +40,7 @@ class RelationalGroupedDataset private[sql] ( groupType: proto.Aggregate.GroupType, pivot: Option[proto.Aggregate.Pivot] = None, groupingSets: Option[Seq[proto.Aggregate.GroupingSets]] = None) - extends api.RelationalGroupedDataset[Dataset] { - type RGD = RelationalGroupedDataset + extends api.RelationalGroupedDataset { import df.sparkSession.RichColumn protected def toDF(aggExprs: Seq[Column]): DataFrame = { diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index aa6258a14b811..04f8eeb5c6d46 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -69,7 +69,7 @@ import org.apache.spark.util.ArrayImplicits._ class SparkSession private[sql] ( private[sql] val client: SparkConnectClient, private val planIdGenerator: AtomicLong) - extends api.SparkSession[Dataset] + extends api.SparkSession with Logging { private[this] val allocator = new RootAllocator() diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala index 11a4a044d20e5..86b1dbe4754e6 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala @@ -20,10 +20,11 @@ package org.apache.spark.sql.catalog import java.util import org.apache.spark.sql.{api, DataFrame, Dataset} +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.types.StructType /** @inheritdoc */ -abstract class Catalog extends api.Catalog[Dataset] { +abstract class Catalog extends api.Catalog { /** @inheritdoc */ override def listDatabases(): Dataset[Database] diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala new file mode 100644 index 0000000000000..7d81f4ead7857 --- /dev/null +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.connect + +import scala.language.implicitConversions + +import org.apache.spark.annotation.DeveloperApi +import org.apache.spark.sql._ + +/** + * Conversions from sql interfaces to the Connect specific implementation. + * + * This class is mainly used by the implementation. In the case of connect it should be extremely + * rare that a developer needs these classes. + * + * We provide both a trait and an object. The trait is useful in situations where an extension + * developer needs to use these conversions in a project covering multiple Spark versions. They + * can create a shim for these conversions, the Spark 4+ version of the shim implements this + * trait, and shims for older versions do not. + */ +@DeveloperApi +trait ConnectConversions { + implicit def castToImpl(session: api.SparkSession): SparkSession = + session.asInstanceOf[SparkSession] + + implicit def castToImpl[T](ds: api.Dataset[T]): Dataset[T] = + ds.asInstanceOf[Dataset[T]] + + implicit def castToImpl(rgds: api.RelationalGroupedDataset): RelationalGroupedDataset = + rgds.asInstanceOf[RelationalGroupedDataset] + + implicit def castToImpl[K, V]( + kvds: api.KeyValueGroupedDataset[K, V]): KeyValueGroupedDataset[K, V] = + kvds.asInstanceOf[KeyValueGroupedDataset[K, V]] +} + +object ConnectConversions extends ConnectConversions diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala index 3b47269875f4a..29fbcc443deb9 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala @@ -26,10 +26,10 @@ import org.apache.spark.connect.proto.ExecutePlanResponse import org.apache.spark.connect.proto.StreamingQueryCommand import org.apache.spark.connect.proto.StreamingQueryCommandResult import org.apache.spark.connect.proto.StreamingQueryManagerCommandResult.StreamingQueryInstance -import org.apache.spark.sql.{api, Dataset, SparkSession} +import org.apache.spark.sql.{api, SparkSession} /** @inheritdoc */ -trait StreamingQuery extends api.StreamingQuery[Dataset] { +trait StreamingQuery extends api.StreamingQuery { /** @inheritdoc */ override def sparkSession: SparkSession diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala index abf03cfbc6722..16f6983efb187 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala @@ -158,6 +158,7 @@ object CheckConnectJvmClientCompatibility { ProblemFilters.exclude[Problem]("org.apache.spark.sql.catalyst.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.columnar.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.connector.*"), + ProblemFilters.exclude[Problem]("org.apache.spark.sql.classic.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.execution.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.internal.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.jdbc.*"), diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index dfe7b14e2ec66..ece4504395f12 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -201,6 +201,8 @@ object MimaExcludes { ProblemFilters.exclude[Problem]("org.apache.spark.sql.execution.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.internal.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.errors.*"), + ProblemFilters.exclude[Problem]("org.apache.spark.sql.classic.*"), + ProblemFilters.exclude[Problem]("org.apache.spark.sql.connect.*"), // DSv2 catalog and expression APIs are unstable yet. We should enable this back. ProblemFilters.exclude[Problem]("org.apache.spark.sql.connector.catalog.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.connector.expressions.*"), diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index 4a8214b2e20a3..d93a52985b772 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -1352,6 +1352,7 @@ trait SharedUnidocSettings { .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/util/kvstore"))) .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/catalyst"))) .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/connect/"))) + .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/classic/"))) .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/execution"))) .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/hive"))) .map(_.filterNot(_.getCanonicalPath.contains("org/apache/spark/sql/catalog/v2/utils"))) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Catalog.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Catalog.scala index fbb665b7f1b1f..a0f51d30dc572 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Catalog.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Catalog.scala @@ -33,7 +33,7 @@ import org.apache.spark.storage.StorageLevel * @since 2.0.0 */ @Stable -abstract class Catalog[DS[U] <: Dataset[U, DS]] { +abstract class Catalog { /** * Returns the current database (namespace) in this session. @@ -54,7 +54,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def listDatabases(): DS[Database] + def listDatabases(): Dataset[Database] /** * Returns a list of databases (namespaces) which name match the specify pattern and available @@ -62,7 +62,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 3.5.0 */ - def listDatabases(pattern: String): DS[Database] + def listDatabases(pattern: String): Dataset[Database] /** * Returns a list of tables/views in the current database (namespace). This includes all @@ -70,7 +70,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def listTables(): DS[Table] + def listTables(): Dataset[Table] /** * Returns a list of tables/views in the specified database (namespace) (the name can be @@ -79,7 +79,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @throws[AnalysisException]("database does not exist") - def listTables(dbName: String): DS[Table] + def listTables(dbName: String): Dataset[Table] /** * Returns a list of tables/views in the specified database (namespace) which name match the @@ -88,7 +88,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 3.5.0 */ @throws[AnalysisException]("database does not exist") - def listTables(dbName: String, pattern: String): DS[Table] + def listTables(dbName: String, pattern: String): Dataset[Table] /** * Returns a list of functions registered in the current database (namespace). This includes all @@ -96,7 +96,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def listFunctions(): DS[Function] + def listFunctions(): Dataset[Function] /** * Returns a list of functions registered in the specified database (namespace) (the name can be @@ -105,7 +105,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @throws[AnalysisException]("database does not exist") - def listFunctions(dbName: String): DS[Function] + def listFunctions(dbName: String): Dataset[Function] /** * Returns a list of functions registered in the specified database (namespace) which name match @@ -115,7 +115,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 3.5.0 */ @throws[AnalysisException]("database does not exist") - def listFunctions(dbName: String, pattern: String): DS[Function] + def listFunctions(dbName: String, pattern: String): Dataset[Function] /** * Returns a list of columns for the given table/view or temporary view. @@ -127,7 +127,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @throws[AnalysisException]("table does not exist") - def listColumns(tableName: String): DS[Column] + def listColumns(tableName: String): Dataset[Column] /** * Returns a list of columns for the given table/view in the specified database under the Hive @@ -143,7 +143,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @throws[AnalysisException]("database or table does not exist") - def listColumns(dbName: String, tableName: String): DS[Column] + def listColumns(dbName: String, tableName: String): Dataset[Column] /** * Get the database (namespace) with the specified name (can be qualified with catalog). This @@ -280,7 +280,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @deprecated("use createTable instead.", "2.2.0") - def createExternalTable(tableName: String, path: String): DS[Row] = { + def createExternalTable(tableName: String, path: String): Dataset[Row] = { createTable(tableName, path) } @@ -293,7 +293,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * identifier is provided, it refers to a table in the current database. * @since 2.2.0 */ - def createTable(tableName: String, path: String): DS[Row] + def createTable(tableName: String, path: String): Dataset[Row] /** * Creates a table from the given path based on a data source and returns the corresponding @@ -305,7 +305,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @deprecated("use createTable instead.", "2.2.0") - def createExternalTable(tableName: String, path: String, source: String): DS[Row] = { + def createExternalTable(tableName: String, path: String, source: String): Dataset[Row] = { createTable(tableName, path, source) } @@ -318,7 +318,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * identifier is provided, it refers to a table in the current database. * @since 2.2.0 */ - def createTable(tableName: String, path: String, source: String): DS[Row] + def createTable(tableName: String, path: String, source: String): Dataset[Row] /** * Creates a table from the given path based on a data source and a set of options. Then, @@ -333,7 +333,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { def createExternalTable( tableName: String, source: String, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable(tableName, source, options) } @@ -349,7 +349,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { def createTable( tableName: String, source: String, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable(tableName, source, options.asScala.toMap) } @@ -366,7 +366,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { def createExternalTable( tableName: String, source: String, - options: Map[String, String]): DS[Row] = { + options: Map[String, String]): Dataset[Row] = { createTable(tableName, source, options) } @@ -379,7 +379,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * identifier is provided, it refers to a table in the current database. * @since 2.2.0 */ - def createTable(tableName: String, source: String, options: Map[String, String]): DS[Row] + def createTable(tableName: String, source: String, options: Map[String, String]): Dataset[Row] /** * Create a table from the given path based on a data source, a schema and a set of options. @@ -395,7 +395,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, schema: StructType, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable(tableName, source, schema, options) } @@ -412,7 +412,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, description: String, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable( tableName, source = source, @@ -433,7 +433,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, description: String, - options: Map[String, String]): DS[Row] + options: Map[String, String]): Dataset[Row] /** * Create a table based on the dataset in a data source, a schema and a set of options. Then, @@ -448,7 +448,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, schema: StructType, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable(tableName, source, schema, options.asScala.toMap) } @@ -466,7 +466,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, schema: StructType, - options: Map[String, String]): DS[Row] = { + options: Map[String, String]): Dataset[Row] = { createTable(tableName, source, schema, options) } @@ -483,7 +483,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { tableName: String, source: String, schema: StructType, - options: Map[String, String]): DS[Row] + options: Map[String, String]): Dataset[Row] /** * Create a table based on the dataset in a data source, a schema and a set of options. Then, @@ -499,7 +499,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { source: String, schema: StructType, description: String, - options: util.Map[String, String]): DS[Row] = { + options: util.Map[String, String]): Dataset[Row] = { createTable( tableName, source = source, @@ -522,7 +522,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { source: String, schema: StructType, description: String, - options: Map[String, String]): DS[Row] + options: Map[String, String]): Dataset[Row] /** * Drops the local temporary view with the given view name in the catalog. If the view has been @@ -670,7 +670,7 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 3.4.0 */ - def listCatalogs(): DS[CatalogMetadata] + def listCatalogs(): Dataset[CatalogMetadata] /** * Returns a list of catalogs which name match the specify pattern and available in this @@ -678,5 +678,5 @@ abstract class Catalog[DS[U] <: Dataset[U, DS]] { * * @since 3.5.0 */ - def listCatalogs(pattern: String): DS[CatalogMetadata] + def listCatalogs(pattern: String): Dataset[CatalogMetadata] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameNaFunctions.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameNaFunctions.scala index 12d3d41aa5546..ef6cc64c058a4 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameNaFunctions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameNaFunctions.scala @@ -30,14 +30,14 @@ import org.apache.spark.util.ArrayImplicits._ * @since 1.3.1 */ @Stable -abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { +abstract class DataFrameNaFunctions { /** * Returns a new `DataFrame` that drops rows containing any null or NaN values. * * @since 1.3.1 */ - def drop(): DS[Row] = drop("any") + def drop(): Dataset[Row] = drop("any") /** * Returns a new `DataFrame` that drops rows containing null or NaN values. @@ -47,7 +47,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(how: String): DS[Row] = drop(toMinNonNulls(how)) + def drop(how: String): Dataset[Row] = drop(toMinNonNulls(how)) /** * Returns a new `DataFrame` that drops rows containing any null or NaN values in the specified @@ -55,7 +55,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(cols: Array[String]): DS[Row] = drop(cols.toImmutableArraySeq) + def drop(cols: Array[String]): Dataset[Row] = drop(cols.toImmutableArraySeq) /** * (Scala-specific) Returns a new `DataFrame` that drops rows containing any null or NaN values @@ -63,7 +63,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(cols: Seq[String]): DS[Row] = drop(cols.size, cols) + def drop(cols: Seq[String]): Dataset[Row] = drop(cols.size, cols) /** * Returns a new `DataFrame` that drops rows containing null or NaN values in the specified @@ -74,7 +74,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(how: String, cols: Array[String]): DS[Row] = drop(how, cols.toImmutableArraySeq) + def drop(how: String, cols: Array[String]): Dataset[Row] = drop(how, cols.toImmutableArraySeq) /** * (Scala-specific) Returns a new `DataFrame` that drops rows containing null or NaN values in @@ -85,7 +85,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(how: String, cols: Seq[String]): DS[Row] = drop(toMinNonNulls(how), cols) + def drop(how: String, cols: Seq[String]): Dataset[Row] = drop(toMinNonNulls(how), cols) /** * Returns a new `DataFrame` that drops rows containing less than `minNonNulls` non-null and @@ -93,7 +93,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(minNonNulls: Int): DS[Row] = drop(Option(minNonNulls)) + def drop(minNonNulls: Int): Dataset[Row] = drop(Option(minNonNulls)) /** * Returns a new `DataFrame` that drops rows containing less than `minNonNulls` non-null and @@ -101,7 +101,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(minNonNulls: Int, cols: Array[String]): DS[Row] = + def drop(minNonNulls: Int, cols: Array[String]): Dataset[Row] = drop(minNonNulls, cols.toImmutableArraySeq) /** @@ -110,7 +110,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def drop(minNonNulls: Int, cols: Seq[String]): DS[Row] = drop(Option(minNonNulls), cols) + def drop(minNonNulls: Int, cols: Seq[String]): Dataset[Row] = drop(Option(minNonNulls), cols) private def toMinNonNulls(how: String): Option[Int] = { how.toLowerCase(util.Locale.ROOT) match { @@ -120,29 +120,29 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { } } - protected def drop(minNonNulls: Option[Int]): DS[Row] + protected def drop(minNonNulls: Option[Int]): Dataset[Row] - protected def drop(minNonNulls: Option[Int], cols: Seq[String]): DS[Row] + protected def drop(minNonNulls: Option[Int], cols: Seq[String]): Dataset[Row] /** * Returns a new `DataFrame` that replaces null or NaN values in numeric columns with `value`. * * @since 2.2.0 */ - def fill(value: Long): DS[Row] + def fill(value: Long): Dataset[Row] /** * Returns a new `DataFrame` that replaces null or NaN values in numeric columns with `value`. * @since 1.3.1 */ - def fill(value: Double): DS[Row] + def fill(value: Double): Dataset[Row] /** * Returns a new `DataFrame` that replaces null values in string columns with `value`. * * @since 1.3.1 */ - def fill(value: String): DS[Row] + def fill(value: String): Dataset[Row] /** * Returns a new `DataFrame` that replaces null or NaN values in specified numeric columns. If a @@ -150,7 +150,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 2.2.0 */ - def fill(value: Long, cols: Array[String]): DS[Row] = fill(value, cols.toImmutableArraySeq) + def fill(value: Long, cols: Array[String]): Dataset[Row] = fill(value, cols.toImmutableArraySeq) /** * Returns a new `DataFrame` that replaces null or NaN values in specified numeric columns. If a @@ -158,7 +158,8 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(value: Double, cols: Array[String]): DS[Row] = fill(value, cols.toImmutableArraySeq) + def fill(value: Double, cols: Array[String]): Dataset[Row] = + fill(value, cols.toImmutableArraySeq) /** * (Scala-specific) Returns a new `DataFrame` that replaces null or NaN values in specified @@ -166,7 +167,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 2.2.0 */ - def fill(value: Long, cols: Seq[String]): DS[Row] + def fill(value: Long, cols: Seq[String]): Dataset[Row] /** * (Scala-specific) Returns a new `DataFrame` that replaces null or NaN values in specified @@ -174,7 +175,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(value: Double, cols: Seq[String]): DS[Row] + def fill(value: Double, cols: Seq[String]): Dataset[Row] /** * Returns a new `DataFrame` that replaces null values in specified string columns. If a @@ -182,7 +183,8 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(value: String, cols: Array[String]): DS[Row] = fill(value, cols.toImmutableArraySeq) + def fill(value: String, cols: Array[String]): Dataset[Row] = + fill(value, cols.toImmutableArraySeq) /** * (Scala-specific) Returns a new `DataFrame` that replaces null values in specified string @@ -190,14 +192,14 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(value: String, cols: Seq[String]): DS[Row] + def fill(value: String, cols: Seq[String]): Dataset[Row] /** * Returns a new `DataFrame` that replaces null values in boolean columns with `value`. * * @since 2.3.0 */ - def fill(value: Boolean): DS[Row] + def fill(value: Boolean): Dataset[Row] /** * (Scala-specific) Returns a new `DataFrame` that replaces null values in specified boolean @@ -205,7 +207,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 2.3.0 */ - def fill(value: Boolean, cols: Seq[String]): DS[Row] + def fill(value: Boolean, cols: Seq[String]): Dataset[Row] /** * Returns a new `DataFrame` that replaces null values in specified boolean columns. If a @@ -213,7 +215,8 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 2.3.0 */ - def fill(value: Boolean, cols: Array[String]): DS[Row] = fill(value, cols.toImmutableArraySeq) + def fill(value: Boolean, cols: Array[String]): Dataset[Row] = + fill(value, cols.toImmutableArraySeq) /** * Returns a new `DataFrame` that replaces null values. @@ -231,7 +234,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(valueMap: util.Map[String, Any]): DS[Row] = fillMap(valueMap.asScala.toSeq) + def fill(valueMap: util.Map[String, Any]): Dataset[Row] = fillMap(valueMap.asScala.toSeq) /** * (Scala-specific) Returns a new `DataFrame` that replaces null values. @@ -251,9 +254,9 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def fill(valueMap: Map[String, Any]): DS[Row] = fillMap(valueMap.toSeq) + def fill(valueMap: Map[String, Any]): Dataset[Row] = fillMap(valueMap.toSeq) - protected def fillMap(values: Seq[(String, Any)]): DS[Row] + protected def fillMap(values: Seq[(String, Any)]): Dataset[Row] /** * Replaces values matching keys in `replacement` map with the corresponding values. @@ -280,7 +283,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def replace[T](col: String, replacement: util.Map[T, T]): DS[Row] = { + def replace[T](col: String, replacement: util.Map[T, T]): Dataset[Row] = { replace[T](col, replacement.asScala.toMap) } @@ -306,7 +309,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def replace[T](cols: Array[String], replacement: util.Map[T, T]): DS[Row] = { + def replace[T](cols: Array[String], replacement: util.Map[T, T]): Dataset[Row] = { replace(cols.toImmutableArraySeq, replacement.asScala.toMap) } @@ -333,7 +336,7 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def replace[T](col: String, replacement: Map[T, T]): DS[Row] + def replace[T](col: String, replacement: Map[T, T]): Dataset[Row] /** * (Scala-specific) Replaces values matching keys in `replacement` map. @@ -355,5 +358,5 @@ abstract class DataFrameNaFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.3.1 */ - def replace[T](cols: Seq[String], replacement: Map[T, T]): DS[Row] + def replace[T](cols: Seq[String], replacement: Map[T, T]): Dataset[Row] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameReader.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameReader.scala index 6e6ab7b9d95a4..c101c52fd0662 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameReader.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameReader.scala @@ -34,7 +34,8 @@ import org.apache.spark.sql.types.StructType * @since 1.4.0 */ @Stable -abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { +abstract class DataFrameReader { + type DS[U] <: Dataset[U] /** * Specifies the input data source format. @@ -149,7 +150,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def load(): DS[Row] + def load(): Dataset[Row] /** * Loads input in as a `DataFrame`, for data sources that require a path (e.g. data backed by a @@ -157,7 +158,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def load(path: String): DS[Row] + def load(path: String): Dataset[Row] /** * Loads input in as a `DataFrame`, for data sources that support multiple paths. Only works if @@ -166,7 +167,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 1.6.0 */ @scala.annotation.varargs - def load(paths: String*): DS[Row] + def load(paths: String*): Dataset[Row] /** * Construct a `DataFrame` representing the database table accessible via JDBC URL url named @@ -179,7 +180,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def jdbc(url: String, table: String, properties: util.Properties): DS[Row] = { + def jdbc(url: String, table: String, properties: util.Properties): Dataset[Row] = { assertNoSpecifiedSchema("jdbc") // properties should override settings in extraOptions. this.extraOptions ++= properties.asScala @@ -223,7 +224,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { lowerBound: Long, upperBound: Long, numPartitions: Int, - connectionProperties: util.Properties): DS[Row] = { + connectionProperties: util.Properties): Dataset[Row] = { // columnName, lowerBound, upperBound and numPartitions override settings in extraOptions. this.extraOptions ++= Map( "partitionColumn" -> columnName, @@ -260,7 +261,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { url: String, table: String, predicates: Array[String], - connectionProperties: util.Properties): DS[Row] + connectionProperties: util.Properties): Dataset[Row] /** * Loads a JSON file and returns the results as a `DataFrame`. @@ -269,7 +270,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def json(path: String): DS[Row] = { + def json(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 json(Seq(path): _*) } @@ -290,7 +291,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @scala.annotation.varargs - def json(paths: String*): DS[Row] = { + def json(paths: String*): Dataset[Row] = { validateJsonSchema() format("json").load(paths: _*) } @@ -306,7 +307,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * input Dataset with one JSON object per record * @since 2.2.0 */ - def json(jsonDataset: DS[String]): DS[Row] + def json(jsonDataset: DS[String]): Dataset[Row] /** * Loads a CSV file and returns the result as a `DataFrame`. See the documentation on the other @@ -314,7 +315,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def csv(path: String): DS[Row] = { + def csv(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 csv(Seq(path): _*) } @@ -340,7 +341,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * input Dataset with one CSV row per record * @since 2.2.0 */ - def csv(csvDataset: DS[String]): DS[Row] + def csv(csvDataset: DS[String]): Dataset[Row] /** * Loads CSV files and returns the result as a `DataFrame`. @@ -356,7 +357,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @scala.annotation.varargs - def csv(paths: String*): DS[Row] = format("csv").load(paths: _*) + def csv(paths: String*): Dataset[Row] = format("csv").load(paths: _*) /** * Loads a XML file and returns the result as a `DataFrame`. See the documentation on the other @@ -364,7 +365,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 4.0.0 */ - def xml(path: String): DS[Row] = { + def xml(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 xml(Seq(path): _*) } @@ -383,7 +384,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 4.0.0 */ @scala.annotation.varargs - def xml(paths: String*): DS[Row] = { + def xml(paths: String*): Dataset[Row] = { validateXmlSchema() format("xml").load(paths: _*) } @@ -398,7 +399,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * input Dataset with one XML object per record * @since 4.0.0 */ - def xml(xmlDataset: DS[String]): DS[Row] + def xml(xmlDataset: DS[String]): Dataset[Row] /** * Loads a Parquet file, returning the result as a `DataFrame`. See the documentation on the @@ -406,7 +407,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def parquet(path: String): DS[Row] = { + def parquet(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 parquet(Seq(path): _*) } @@ -421,7 +422,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 1.4.0 */ @scala.annotation.varargs - def parquet(paths: String*): DS[Row] = format("parquet").load(paths: _*) + def parquet(paths: String*): Dataset[Row] = format("parquet").load(paths: _*) /** * Loads an ORC file and returns the result as a `DataFrame`. @@ -430,7 +431,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * input path * @since 1.5.0 */ - def orc(path: String): DS[Row] = { + def orc(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 orc(Seq(path): _*) } @@ -447,7 +448,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @scala.annotation.varargs - def orc(paths: String*): DS[Row] = format("orc").load(paths: _*) + def orc(paths: String*): Dataset[Row] = format("orc").load(paths: _*) /** * Returns the specified table/view as a `DataFrame`. If it's a table, it must support batch @@ -462,7 +463,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * database. Note that, the global temporary view database is also valid here. * @since 1.4.0 */ - def table(tableName: String): DS[Row] + def table(tableName: String): Dataset[Row] /** * Loads text files and returns a `DataFrame` whose schema starts with a string column named @@ -471,7 +472,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def text(path: String): DS[Row] = { + def text(path: String): Dataset[Row] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 text(Seq(path): _*) } @@ -499,14 +500,14 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 1.6.0 */ @scala.annotation.varargs - def text(paths: String*): DS[Row] = format("text").load(paths: _*) + def text(paths: String*): Dataset[Row] = format("text").load(paths: _*) /** * Loads text files and returns a [[Dataset]] of String. See the documentation on the other * overloaded `textFile()` method for more details. * @since 2.0.0 */ - def textFile(path: String): DS[String] = { + def textFile(path: String): Dataset[String] = { // This method ensures that calls that explicit need single argument works, see SPARK-16009 textFile(Seq(path): _*) } @@ -534,7 +535,7 @@ abstract class DataFrameReader[DS[U] <: Dataset[U, DS]] { * @since 2.0.0 */ @scala.annotation.varargs - def textFile(paths: String*): DS[String] = { + def textFile(paths: String*): Dataset[String] = { assertNoSpecifiedSchema("textFile") text(paths: _*).select("value").as(StringEncoder) } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameStatFunctions.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameStatFunctions.scala index fc1680231be5b..ae7c256b30ace 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameStatFunctions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/DataFrameStatFunctions.scala @@ -34,8 +34,8 @@ import org.apache.spark.util.sketch.{BloomFilter, CountMinSketch} * @since 1.4.0 */ @Stable -abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { - protected def df: DS[Row] +abstract class DataFrameStatFunctions { + protected def df: Dataset[Row] /** * Calculates the approximate quantiles of a numerical column of a DataFrame. @@ -202,7 +202,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def crosstab(col1: String, col2: String): DS[Row] + def crosstab(col1: String, col2: String): Dataset[Row] /** * Finding frequent items for columns, possibly with false positives. Using the frequent element @@ -246,7 +246,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * }}} * @since 1.4.0 */ - def freqItems(cols: Array[String], support: Double): DS[Row] = + def freqItems(cols: Array[String], support: Double): Dataset[Row] = freqItems(cols.toImmutableArraySeq, support) /** @@ -263,7 +263,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * A Local DataFrame with the Array of frequent items for each column. * @since 1.4.0 */ - def freqItems(cols: Array[String]): DS[Row] = freqItems(cols, 0.01) + def freqItems(cols: Array[String]): Dataset[Row] = freqItems(cols, 0.01) /** * (Scala-specific) Finding frequent items for columns, possibly with false positives. Using the @@ -307,7 +307,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.4.0 */ - def freqItems(cols: Seq[String], support: Double): DS[Row] + def freqItems(cols: Seq[String], support: Double): Dataset[Row] /** * (Scala-specific) Finding frequent items for columns, possibly with false positives. Using the @@ -324,7 +324,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * A Local DataFrame with the Array of frequent items for each column. * @since 1.4.0 */ - def freqItems(cols: Seq[String]): DS[Row] = freqItems(cols, 0.01) + def freqItems(cols: Seq[String]): Dataset[Row] = freqItems(cols, 0.01) /** * Returns a stratified sample without replacement based on the fraction given on each stratum. @@ -356,7 +356,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.5.0 */ - def sampleBy[T](col: String, fractions: Map[T, Double], seed: Long): DS[Row] = { + def sampleBy[T](col: String, fractions: Map[T, Double], seed: Long): Dataset[Row] = { sampleBy(Column(col), fractions, seed) } @@ -376,7 +376,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * * @since 1.5.0 */ - def sampleBy[T](col: String, fractions: ju.Map[T, jl.Double], seed: Long): DS[Row] = { + def sampleBy[T](col: String, fractions: ju.Map[T, jl.Double], seed: Long): Dataset[Row] = { sampleBy(col, fractions.asScala.toMap.asInstanceOf[Map[T, Double]], seed) } @@ -413,7 +413,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * * @since 3.0.0 */ - def sampleBy[T](col: Column, fractions: Map[T, Double], seed: Long): DS[Row] + def sampleBy[T](col: Column, fractions: Map[T, Double], seed: Long): Dataset[Row] /** * (Java-specific) Returns a stratified sample without replacement based on the fraction given @@ -432,7 +432,7 @@ abstract class DataFrameStatFunctions[DS[U] <: Dataset[U, DS]] { * a new `DataFrame` that represents the stratified sample * @since 3.0.0 */ - def sampleBy[T](col: Column, fractions: ju.Map[T, jl.Double], seed: Long): DS[Row] = { + def sampleBy[T](col: Column, fractions: ju.Map[T, jl.Double], seed: Long): Dataset[Row] = { sampleBy(col, fractions.asScala.toMap.asInstanceOf[Map[T, Double]], seed) } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala index fb8b6f2f483a1..284a69fe6ee3e 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala @@ -119,10 +119,10 @@ import org.apache.spark.util.SparkClassUtils * @since 1.6.0 */ @Stable -abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { - type RGD <: RelationalGroupedDataset[DS] +abstract class Dataset[T] extends Serializable { + type DS[U] <: Dataset[U] - def sparkSession: SparkSession[DS] + def sparkSession: SparkSession val encoder: Encoder[T] @@ -136,7 +136,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { */ // This is declared with parentheses to prevent the Scala compiler from treating // `ds.toDF("1")` as invoking this toDF and then apply on the returned DataFrame. - def toDF(): DS[Row] + def toDF(): Dataset[Row] /** * Returns a new Dataset where each record has been mapped on to the specified type. The method @@ -157,7 +157,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def as[U: Encoder]: DS[U] + def as[U: Encoder]: Dataset[U] /** * Returns a new DataFrame where each row is reconciled to match the specified schema. Spark @@ -175,7 +175,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 3.4.0 */ - def to(schema: StructType): DS[Row] + def to(schema: StructType): Dataset[Row] /** * Converts this strongly typed collection of data to generic `DataFrame` with columns renamed. @@ -191,7 +191,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def toDF(colNames: String*): DS[Row] + def toDF(colNames: String*): Dataset[Row] /** * Returns the schema of this Dataset. @@ -312,7 +312,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 2.1.0 */ - def checkpoint(): DS[T] = checkpoint(eager = true, reliableCheckpoint = true) + def checkpoint(): Dataset[T] = checkpoint(eager = true, reliableCheckpoint = true) /** * Returns a checkpointed version of this Dataset. Checkpointing can be used to truncate the @@ -331,7 +331,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 2.1.0 */ - def checkpoint(eager: Boolean): DS[T] = checkpoint(eager = eager, reliableCheckpoint = true) + def checkpoint(eager: Boolean): Dataset[T] = + checkpoint(eager = eager, reliableCheckpoint = true) /** * Eagerly locally checkpoints a Dataset and return the new Dataset. Checkpointing can be used @@ -342,7 +343,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 2.3.0 */ - def localCheckpoint(): DS[T] = checkpoint(eager = true, reliableCheckpoint = false) + def localCheckpoint(): Dataset[T] = checkpoint(eager = true, reliableCheckpoint = false) /** * Locally checkpoints a Dataset and return the new Dataset. Checkpointing can be used to @@ -361,7 +362,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 2.3.0 */ - def localCheckpoint(eager: Boolean): DS[T] = + def localCheckpoint(eager: Boolean): Dataset[T] = checkpoint(eager = eager, reliableCheckpoint = false) /** @@ -373,7 +374,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * Whether to create a reliable checkpoint saved to files inside the checkpoint directory. If * false creates a local checkpoint using the caching subsystem */ - protected def checkpoint(eager: Boolean, reliableCheckpoint: Boolean): DS[T] + protected def checkpoint(eager: Boolean, reliableCheckpoint: Boolean): Dataset[T] /** * Defines an event time watermark for this [[Dataset]]. A watermark tracks a point in time @@ -400,7 +401,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { */ // We only accept an existing column name, not a derived column here as a watermark that is // defined on a derived column cannot referenced elsewhere in the plan. - def withWatermark(eventTime: String, delayThreshold: String): DS[T] + def withWatermark(eventTime: String, delayThreshold: String): Dataset[T] /** * Displays the Dataset in a tabular form. Strings more than 20 characters will be truncated, @@ -551,7 +552,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 1.6.0 */ - def na: DataFrameNaFunctions[DS] + def na: DataFrameNaFunctions /** * Returns a [[DataFrameStatFunctions]] for working statistic functions support. @@ -563,7 +564,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 1.6.0 */ - def stat: DataFrameStatFunctions[DS] + def stat: DataFrameStatFunctions /** * Join with another `DataFrame`. @@ -575,7 +576,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_]): DS[Row] + def join(right: DS[_]): Dataset[Row] /** * Inner equi-join with another `DataFrame` using the given column. @@ -601,7 +602,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_], usingColumn: String): DS[Row] = { + def join(right: DS[_], usingColumn: String): Dataset[Row] = { join(right, Seq(usingColumn)) } @@ -617,7 +618,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def join(right: DS[_], usingColumns: Array[String]): DS[Row] = { + def join(right: DS[_], usingColumns: Array[String]): Dataset[Row] = { join(right, usingColumns.toImmutableArraySeq) } @@ -645,7 +646,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_], usingColumns: Seq[String]): DS[Row] = { + def join(right: DS[_], usingColumns: Seq[String]): Dataset[Row] = { join(right, usingColumns, "inner") } @@ -675,7 +676,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def join(right: DS[_], usingColumn: String, joinType: String): DS[Row] = { + def join(right: DS[_], usingColumn: String, joinType: String): Dataset[Row] = { join(right, Seq(usingColumn), joinType) } @@ -696,7 +697,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def join(right: DS[_], usingColumns: Array[String], joinType: String): DS[Row] = { + def join(right: DS[_], usingColumns: Array[String], joinType: String): Dataset[Row] = { join(right, usingColumns.toImmutableArraySeq, joinType) } @@ -726,7 +727,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_], usingColumns: Seq[String], joinType: String): DS[Row] + def join(right: DS[_], usingColumns: Seq[String], joinType: String): Dataset[Row] /** * Inner join with another `DataFrame`, using the given join expression. @@ -740,7 +741,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_], joinExprs: Column): DS[Row] = + def join(right: DS[_], joinExprs: Column): Dataset[Row] = join(right, joinExprs, "inner") /** @@ -770,7 +771,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def join(right: DS[_], joinExprs: Column, joinType: String): DS[Row] + def join(right: DS[_], joinExprs: Column, joinType: String): Dataset[Row] /** * Explicit cartesian join with another `DataFrame`. @@ -782,7 +783,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.1.0 */ - def crossJoin(right: DS[_]): DS[Row] + def crossJoin(right: DS[_]): Dataset[Row] /** * Joins this Dataset returning a `Tuple2` for each pair where `condition` evaluates to true. @@ -806,7 +807,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def joinWith[U](other: DS[U], condition: Column, joinType: String): DS[(T, U)] + def joinWith[U](other: DS[U], condition: Column, joinType: String): Dataset[(T, U)] /** * Using inner equi-join to join this Dataset returning a `Tuple2` for each pair where @@ -819,11 +820,11 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def joinWith[U](other: DS[U], condition: Column): DS[(T, U)] = { + def joinWith[U](other: DS[U], condition: Column): Dataset[(T, U)] = { joinWith(other, condition, "inner") } - protected def sortInternal(global: Boolean, sortExprs: Seq[Column]): DS[T] + protected def sortInternal(global: Boolean, sortExprs: Seq[Column]): Dataset[T] /** * Returns a new Dataset with each partition sorted by the given expressions. @@ -834,7 +835,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def sortWithinPartitions(sortCol: String, sortCols: String*): DS[T] = { + def sortWithinPartitions(sortCol: String, sortCols: String*): Dataset[T] = { sortWithinPartitions((sortCol +: sortCols).map(Column(_)): _*) } @@ -847,7 +848,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def sortWithinPartitions(sortExprs: Column*): DS[T] = { + def sortWithinPartitions(sortExprs: Column*): Dataset[T] = { sortInternal(global = false, sortExprs) } @@ -864,7 +865,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def sort(sortCol: String, sortCols: String*): DS[T] = { + def sort(sortCol: String, sortCols: String*): Dataset[T] = { sort((sortCol +: sortCols).map(Column(_)): _*) } @@ -878,7 +879,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def sort(sortExprs: Column*): DS[T] = { + def sort(sortExprs: Column*): Dataset[T] = { sortInternal(global = true, sortExprs) } @@ -890,7 +891,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def orderBy(sortCol: String, sortCols: String*): DS[T] = sort(sortCol, sortCols: _*) + def orderBy(sortCol: String, sortCols: String*): Dataset[T] = sort(sortCol, sortCols: _*) /** * Returns a new Dataset sorted by the given expressions. This is an alias of the `sort` @@ -900,7 +901,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def orderBy(sortExprs: Column*): DS[T] = sort(sortExprs: _*) + def orderBy(sortExprs: Column*): Dataset[T] = sort(sortExprs: _*) /** * Specifies some hint on the current Dataset. As an example, the following code specifies that @@ -926,7 +927,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.2.0 */ @scala.annotation.varargs - def hint(name: String, parameters: Any*): DS[T] + def hint(name: String, parameters: Any*): Dataset[T] /** * Selects column based on the column name and returns it as a [[org.apache.spark.sql.Column]]. @@ -975,7 +976,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def as(alias: String): DS[T] + def as(alias: String): Dataset[T] /** * (Scala-specific) Returns a new Dataset with an alias set. @@ -983,7 +984,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def as(alias: Symbol): DS[T] = as(alias.name) + def as(alias: Symbol): Dataset[T] = as(alias.name) /** * Returns a new Dataset with an alias set. Same as `as`. @@ -991,7 +992,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def alias(alias: String): DS[T] = as(alias) + def alias(alias: String): Dataset[T] = as(alias) /** * (Scala-specific) Returns a new Dataset with an alias set. Same as `as`. @@ -999,7 +1000,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def alias(alias: Symbol): DS[T] = as(alias) + def alias(alias: Symbol): Dataset[T] = as(alias) /** * Selects a set of column based expressions. @@ -1011,7 +1012,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def select(cols: Column*): DS[Row] + def select(cols: Column*): Dataset[Row] /** * Selects a set of columns. This is a variant of `select` that can only select existing columns @@ -1027,7 +1028,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def select(col: String, cols: String*): DS[Row] = select((col +: cols).map(Column(_)): _*) + def select(col: String, cols: String*): Dataset[Row] = select((col +: cols).map(Column(_)): _*) /** * Selects a set of SQL expressions. This is a variant of `select` that accepts SQL expressions. @@ -1042,7 +1043,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def selectExpr(exprs: String*): DS[Row] = select(exprs.map(functions.expr): _*) + def selectExpr(exprs: String*): Dataset[Row] = select(exprs.map(functions.expr): _*) /** * Returns a new Dataset by computing the given [[org.apache.spark.sql.Column]] expression for @@ -1056,14 +1057,14 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def select[U1](c1: TypedColumn[T, U1]): DS[U1] + def select[U1](c1: TypedColumn[T, U1]): Dataset[U1] /** * Internal helper function for building typed selects that return tuples. For simplicity and * code reuse, we do this without the help of the type system and then use helper functions that * cast appropriately for the user facing interface. */ - protected def selectUntyped(columns: TypedColumn[_, _]*): DS[_] + protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] /** * Returns a new Dataset by computing the given [[org.apache.spark.sql.Column]] expressions for @@ -1072,8 +1073,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def select[U1, U2](c1: TypedColumn[T, U1], c2: TypedColumn[T, U2]): DS[(U1, U2)] = - selectUntyped(c1, c2).asInstanceOf[DS[(U1, U2)]] + def select[U1, U2](c1: TypedColumn[T, U1], c2: TypedColumn[T, U2]): Dataset[(U1, U2)] = + selectUntyped(c1, c2).asInstanceOf[Dataset[(U1, U2)]] /** * Returns a new Dataset by computing the given [[org.apache.spark.sql.Column]] expressions for @@ -1085,8 +1086,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { def select[U1, U2, U3]( c1: TypedColumn[T, U1], c2: TypedColumn[T, U2], - c3: TypedColumn[T, U3]): DS[(U1, U2, U3)] = - selectUntyped(c1, c2, c3).asInstanceOf[DS[(U1, U2, U3)]] + c3: TypedColumn[T, U3]): Dataset[(U1, U2, U3)] = + selectUntyped(c1, c2, c3).asInstanceOf[Dataset[(U1, U2, U3)]] /** * Returns a new Dataset by computing the given [[org.apache.spark.sql.Column]] expressions for @@ -1099,8 +1100,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { c1: TypedColumn[T, U1], c2: TypedColumn[T, U2], c3: TypedColumn[T, U3], - c4: TypedColumn[T, U4]): DS[(U1, U2, U3, U4)] = - selectUntyped(c1, c2, c3, c4).asInstanceOf[DS[(U1, U2, U3, U4)]] + c4: TypedColumn[T, U4]): Dataset[(U1, U2, U3, U4)] = + selectUntyped(c1, c2, c3, c4).asInstanceOf[Dataset[(U1, U2, U3, U4)]] /** * Returns a new Dataset by computing the given [[org.apache.spark.sql.Column]] expressions for @@ -1114,8 +1115,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { c2: TypedColumn[T, U2], c3: TypedColumn[T, U3], c4: TypedColumn[T, U4], - c5: TypedColumn[T, U5]): DS[(U1, U2, U3, U4, U5)] = - selectUntyped(c1, c2, c3, c4, c5).asInstanceOf[DS[(U1, U2, U3, U4, U5)]] + c5: TypedColumn[T, U5]): Dataset[(U1, U2, U3, U4, U5)] = + selectUntyped(c1, c2, c3, c4, c5).asInstanceOf[Dataset[(U1, U2, U3, U4, U5)]] /** * Filters rows using the given condition. @@ -1128,7 +1129,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def filter(condition: Column): DS[T] + def filter(condition: Column): Dataset[T] /** * Filters rows using the given SQL expression. @@ -1139,7 +1140,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def filter(conditionExpr: String): DS[T] = + def filter(conditionExpr: String): Dataset[T] = filter(functions.expr(conditionExpr)) /** @@ -1149,7 +1150,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def filter(func: T => Boolean): DS[T] + def filter(func: T => Boolean): Dataset[T] /** * (Java-specific) Returns a new Dataset that only contains elements where `func` returns @@ -1158,7 +1159,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def filter(func: FilterFunction[T]): DS[T] + def filter(func: FilterFunction[T]): Dataset[T] /** * Filters rows using the given condition. This is an alias for `filter`. @@ -1171,7 +1172,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def where(condition: Column): DS[T] = filter(condition) + def where(condition: Column): Dataset[T] = filter(condition) /** * Filters rows using the given SQL expression. @@ -1182,7 +1183,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def where(conditionExpr: String): DS[T] = filter(conditionExpr) + def where(conditionExpr: String): Dataset[T] = filter(conditionExpr) /** * Groups the Dataset using the specified columns, so we can run aggregation on them. See @@ -1203,7 +1204,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def groupBy(cols: Column*): RGD + def groupBy(cols: Column*): RelationalGroupedDataset /** * Groups the Dataset using the specified columns, so that we can run aggregation on them. See @@ -1227,7 +1228,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def groupBy(col1: String, cols: String*): RGD = groupBy((col1 +: cols).map(col): _*) + def groupBy(col1: String, cols: String*): RelationalGroupedDataset = groupBy( + (col1 +: cols).map(col): _*) /** * Create a multi-dimensional rollup for the current Dataset using the specified columns, so we @@ -1249,7 +1251,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def rollup(cols: Column*): RGD + def rollup(cols: Column*): RelationalGroupedDataset /** * Create a multi-dimensional rollup for the current Dataset using the specified columns, so we @@ -1274,7 +1276,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def rollup(col1: String, cols: String*): RGD = rollup((col1 +: cols).map(col): _*) + def rollup(col1: String, cols: String*): RelationalGroupedDataset = rollup( + (col1 +: cols).map(col): _*) /** * Create a multi-dimensional cube for the current Dataset using the specified columns, so we @@ -1296,7 +1299,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def cube(cols: Column*): RGD + def cube(cols: Column*): RelationalGroupedDataset /** * Create a multi-dimensional cube for the current Dataset using the specified columns, so we @@ -1321,7 +1324,8 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def cube(col1: String, cols: String*): RGD = cube((col1 +: cols).map(col): _*) + def cube(col1: String, cols: String*): RelationalGroupedDataset = cube( + (col1 +: cols).map(col): _*) /** * Create multi-dimensional aggregation for the current Dataset using the specified grouping @@ -1343,7 +1347,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 4.0.0 */ @scala.annotation.varargs - def groupingSets(groupingSets: Seq[Seq[Column]], cols: Column*): RGD + def groupingSets(groupingSets: Seq[Seq[Column]], cols: Column*): RelationalGroupedDataset /** * (Scala-specific) Aggregates on the entire Dataset without groups. @@ -1356,7 +1360,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def agg(aggExpr: (String, String), aggExprs: (String, String)*): DS[Row] = { + def agg(aggExpr: (String, String), aggExprs: (String, String)*): Dataset[Row] = { groupBy().agg(aggExpr, aggExprs: _*) } @@ -1371,7 +1375,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def agg(exprs: Map[String, String]): DS[Row] = groupBy().agg(exprs) + def agg(exprs: Map[String, String]): Dataset[Row] = groupBy().agg(exprs) /** * (Java-specific) Aggregates on the entire Dataset without groups. @@ -1384,7 +1388,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def agg(exprs: util.Map[String, String]): DS[Row] = groupBy().agg(exprs) + def agg(exprs: util.Map[String, String]): Dataset[Row] = groupBy().agg(exprs) /** * Aggregates on the entire Dataset without groups. @@ -1398,7 +1402,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def agg(expr: Column, exprs: Column*): DS[Row] = groupBy().agg(expr, exprs: _*) + def agg(expr: Column, exprs: Column*): Dataset[Row] = groupBy().agg(expr, exprs: _*) /** * (Scala-specific) Reduces the elements of this Dataset using the specified binary function. @@ -1479,7 +1483,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { ids: Array[Column], values: Array[Column], variableColumnName: String, - valueColumnName: String): DS[Row] + valueColumnName: String): Dataset[Row] /** * Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns @@ -1502,7 +1506,10 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def unpivot(ids: Array[Column], variableColumnName: String, valueColumnName: String): DS[Row] + def unpivot( + ids: Array[Column], + variableColumnName: String, + valueColumnName: String): Dataset[Row] /** * Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns @@ -1526,7 +1533,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { ids: Array[Column], values: Array[Column], variableColumnName: String, - valueColumnName: String): DS[Row] = + valueColumnName: String): Dataset[Row] = unpivot(ids, values, variableColumnName, valueColumnName) /** @@ -1548,7 +1555,10 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def melt(ids: Array[Column], variableColumnName: String, valueColumnName: String): DS[Row] = + def melt( + ids: Array[Column], + variableColumnName: String, + valueColumnName: String): Dataset[Row] = unpivot(ids, variableColumnName, valueColumnName) /** @@ -1611,7 +1621,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 4.0.0 */ - def transpose(indexColumn: Column): DS[Row] + def transpose(indexColumn: Column): Dataset[Row] /** * Transposes a DataFrame, switching rows to columns. This function transforms the DataFrame @@ -1630,7 +1640,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 4.0.0 */ - def transpose(): DS[Row] + def transpose(): Dataset[Row] /** * Define (named) metrics to observe on the Dataset. This method returns an 'observed' Dataset @@ -1651,7 +1661,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 3.0.0 */ @scala.annotation.varargs - def observe(name: String, expr: Column, exprs: Column*): DS[T] + def observe(name: String, expr: Column, exprs: Column*): Dataset[T] /** * Observe (named) metrics through an `org.apache.spark.sql.Observation` instance. This method @@ -1674,7 +1684,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 3.3.0 */ @scala.annotation.varargs - def observe(observation: Observation, expr: Column, exprs: Column*): DS[T] + def observe(observation: Observation, expr: Column, exprs: Column*): Dataset[T] /** * Returns a new Dataset by taking the first `n` rows. The difference between this function and @@ -1684,7 +1694,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def limit(n: Int): DS[T] + def limit(n: Int): Dataset[T] /** * Returns a new Dataset by skipping the first `n` rows. @@ -1692,7 +1702,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 3.4.0 */ - def offset(n: Int): DS[T] + def offset(n: Int): Dataset[T] /** * Returns a new Dataset containing union of rows in this Dataset and another Dataset. @@ -1724,7 +1734,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def union(other: DS[T]): DS[T] + def union(other: DS[T]): Dataset[T] /** * Returns a new Dataset containing union of rows in this Dataset and another Dataset. This is @@ -1738,7 +1748,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def unionAll(other: DS[T]): DS[T] = union(other) + def unionAll(other: DS[T]): Dataset[T] = union(other) /** * Returns a new Dataset containing union of rows in this Dataset and another Dataset. @@ -1769,7 +1779,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.3.0 */ - def unionByName(other: DS[T]): DS[T] = unionByName(other, allowMissingColumns = false) + def unionByName(other: DS[T]): Dataset[T] = unionByName(other, allowMissingColumns = false) /** * Returns a new Dataset containing union of rows in this Dataset and another Dataset. @@ -1813,7 +1823,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 3.1.0 */ - def unionByName(other: DS[T], allowMissingColumns: Boolean): DS[T] + def unionByName(other: DS[T], allowMissingColumns: Boolean): Dataset[T] /** * Returns a new Dataset containing rows only in both this Dataset and another Dataset. This is @@ -1825,7 +1835,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def intersect(other: DS[T]): DS[T] + def intersect(other: DS[T]): Dataset[T] /** * Returns a new Dataset containing rows only in both this Dataset and another Dataset while @@ -1838,7 +1848,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.4.0 */ - def intersectAll(other: DS[T]): DS[T] + def intersectAll(other: DS[T]): Dataset[T] /** * Returns a new Dataset containing rows in this Dataset but not in another Dataset. This is @@ -1850,7 +1860,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def except(other: DS[T]): DS[T] + def except(other: DS[T]): Dataset[T] /** * Returns a new Dataset containing rows in this Dataset but not in another Dataset while @@ -1863,7 +1873,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.4.0 */ - def exceptAll(other: DS[T]): DS[T] + def exceptAll(other: DS[T]): Dataset[T] /** * Returns a new [[Dataset]] by sampling a fraction of rows (without replacement), using a @@ -1879,7 +1889,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.3.0 */ - def sample(fraction: Double, seed: Long): DS[T] = { + def sample(fraction: Double, seed: Long): Dataset[T] = { sample(withReplacement = false, fraction = fraction, seed = seed) } @@ -1895,7 +1905,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.3.0 */ - def sample(fraction: Double): DS[T] = { + def sample(fraction: Double): Dataset[T] = { sample(withReplacement = false, fraction = fraction) } @@ -1914,7 +1924,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def sample(withReplacement: Boolean, fraction: Double, seed: Long): DS[T] + def sample(withReplacement: Boolean, fraction: Double, seed: Long): Dataset[T] /** * Returns a new [[Dataset]] by sampling a fraction of rows, using a random seed. @@ -1931,7 +1941,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def sample(withReplacement: Boolean, fraction: Double): DS[T] = { + def sample(withReplacement: Boolean, fraction: Double): Dataset[T] = { sample(withReplacement, fraction, SparkClassUtils.random.nextLong) } @@ -1948,7 +1958,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def randomSplit(weights: Array[Double], seed: Long): Array[_ <: DS[T]] + def randomSplit(weights: Array[Double], seed: Long): Array[_ <: Dataset[T]] /** * Returns a Java list that contains randomly split Dataset with the provided weights. @@ -1960,7 +1970,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def randomSplitAsList(weights: Array[Double], seed: Long): util.List[_ <: DS[T]] + def randomSplitAsList(weights: Array[Double], seed: Long): util.List[_ <: Dataset[T]] /** * Randomly splits this Dataset with the provided weights. @@ -1970,7 +1980,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def randomSplit(weights: Array[Double]): Array[_ <: DS[T]] + def randomSplit(weights: Array[Double]): Array[_ <: Dataset[T]] /** * (Scala-specific) Returns a new Dataset where each row has been expanded to zero or more rows @@ -1983,7 +1993,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * * {{{ * case class Book(title: String, words: String) - * val ds: DS[Book] + * val ds: Dataset[Book] * * val allWords = ds.select($"title", explode(split($"words", " ")).as("word")) * @@ -2000,7 +2010,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @deprecated("use flatMap() or select() with functions.explode() instead", "2.0.0") - def explode[A <: Product: TypeTag](input: Column*)(f: Row => IterableOnce[A]): DS[Row] + def explode[A <: Product: TypeTag](input: Column*)(f: Row => IterableOnce[A]): Dataset[Row] /** * (Scala-specific) Returns a new Dataset where a single column has been expanded to zero or @@ -2026,7 +2036,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { */ @deprecated("use flatMap() or select() with functions.explode() instead", "2.0.0") def explode[A, B: TypeTag](inputColumn: String, outputColumn: String)( - f: A => IterableOnce[B]): DS[Row] + f: A => IterableOnce[B]): Dataset[Row] /** * Returns a new Dataset by adding a column or replacing the existing column that has the same @@ -2043,7 +2053,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def withColumn(colName: String, col: Column): DS[Row] = withColumns(Seq(colName), Seq(col)) + def withColumn(colName: String, col: Column): Dataset[Row] = withColumns(Seq(colName), Seq(col)) /** * (Scala-specific) Returns a new Dataset by adding columns or replacing the existing columns @@ -2055,7 +2065,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.3.0 */ - def withColumns(colsMap: Map[String, Column]): DS[Row] = { + def withColumns(colsMap: Map[String, Column]): Dataset[Row] = { val (colNames, newCols) = colsMap.toSeq.unzip withColumns(colNames, newCols) } @@ -2070,13 +2080,14 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.3.0 */ - def withColumns(colsMap: util.Map[String, Column]): DS[Row] = withColumns(colsMap.asScala.toMap) + def withColumns(colsMap: util.Map[String, Column]): Dataset[Row] = withColumns( + colsMap.asScala.toMap) /** * Returns a new Dataset by adding columns or replacing the existing columns that has the same * names. */ - protected def withColumns(colNames: Seq[String], cols: Seq[Column]): DS[Row] + protected def withColumns(colNames: Seq[String], cols: Seq[Column]): Dataset[Row] /** * Returns a new Dataset with a column renamed. This is a no-op if schema doesn't contain @@ -2085,7 +2096,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def withColumnRenamed(existingName: String, newName: String): DS[Row] = + def withColumnRenamed(existingName: String, newName: String): Dataset[Row] = withColumnsRenamed(Seq(existingName), Seq(newName)) /** @@ -2100,7 +2111,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 3.4.0 */ @throws[AnalysisException] - def withColumnsRenamed(colsMap: Map[String, String]): DS[Row] = { + def withColumnsRenamed(colsMap: Map[String, String]): Dataset[Row] = { val (colNames, newColNames) = colsMap.toSeq.unzip withColumnsRenamed(colNames, newColNames) } @@ -2114,10 +2125,10 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.4.0 */ - def withColumnsRenamed(colsMap: util.Map[String, String]): DS[Row] = + def withColumnsRenamed(colsMap: util.Map[String, String]): Dataset[Row] = withColumnsRenamed(colsMap.asScala.toMap) - protected def withColumnsRenamed(colNames: Seq[String], newColNames: Seq[String]): DS[Row] + protected def withColumnsRenamed(colNames: Seq[String], newColNames: Seq[String]): Dataset[Row] /** * Returns a new Dataset by updating an existing column with metadata. @@ -2125,7 +2136,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 3.3.0 */ - def withMetadata(columnName: String, metadata: Metadata): DS[Row] + def withMetadata(columnName: String, metadata: Metadata): Dataset[Row] /** * Returns a new Dataset with a column dropped. This is a no-op if schema doesn't contain column @@ -2198,7 +2209,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def drop(colName: String): DS[Row] = drop(colName :: Nil: _*) + def drop(colName: String): Dataset[Row] = drop(colName :: Nil: _*) /** * Returns a new Dataset with columns dropped. This is a no-op if schema doesn't contain column @@ -2211,7 +2222,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def drop(colNames: String*): DS[Row] + def drop(colNames: String*): Dataset[Row] /** * Returns a new Dataset with column dropped. @@ -2226,7 +2237,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group untypedrel * @since 2.0.0 */ - def drop(col: Column): DS[Row] = drop(col, Nil: _*) + def drop(col: Column): Dataset[Row] = drop(col, Nil: _*) /** * Returns a new Dataset with columns dropped. @@ -2238,7 +2249,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 3.4.0 */ @scala.annotation.varargs - def drop(col: Column, cols: Column*): DS[Row] + def drop(col: Column, cols: Column*): Dataset[Row] /** * Returns a new Dataset that contains only the unique rows from this Dataset. This is an alias @@ -2253,7 +2264,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def dropDuplicates(): DS[T] + def dropDuplicates(): Dataset[T] /** * (Scala-specific) Returns a new Dataset with duplicate rows removed, considering only the @@ -2268,7 +2279,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def dropDuplicates(colNames: Seq[String]): DS[T] + def dropDuplicates(colNames: Seq[String]): Dataset[T] /** * Returns a new Dataset with duplicate rows removed, considering only the subset of columns. @@ -2282,7 +2293,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def dropDuplicates(colNames: Array[String]): DS[T] = + def dropDuplicates(colNames: Array[String]): Dataset[T] = dropDuplicates(colNames.toImmutableArraySeq) /** @@ -2299,7 +2310,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def dropDuplicates(col1: String, cols: String*): DS[T] = { + def dropDuplicates(col1: String, cols: String*): Dataset[T] = { val colNames: Seq[String] = col1 +: cols dropDuplicates(colNames) } @@ -2321,7 +2332,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 3.5.0 */ - def dropDuplicatesWithinWatermark(): DS[T] + def dropDuplicatesWithinWatermark(): Dataset[T] /** * Returns a new Dataset with duplicates rows removed, considering only the subset of columns, @@ -2341,7 +2352,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 3.5.0 */ - def dropDuplicatesWithinWatermark(colNames: Seq[String]): DS[T] + def dropDuplicatesWithinWatermark(colNames: Seq[String]): Dataset[T] /** * Returns a new Dataset with duplicates rows removed, considering only the subset of columns, @@ -2361,7 +2372,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 3.5.0 */ - def dropDuplicatesWithinWatermark(colNames: Array[String]): DS[T] = { + def dropDuplicatesWithinWatermark(colNames: Array[String]): Dataset[T] = { dropDuplicatesWithinWatermark(colNames.toImmutableArraySeq) } @@ -2384,7 +2395,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 3.5.0 */ @scala.annotation.varargs - def dropDuplicatesWithinWatermark(col1: String, cols: String*): DS[T] = { + def dropDuplicatesWithinWatermark(col1: String, cols: String*): Dataset[T] = { val colNames: Seq[String] = col1 +: cols dropDuplicatesWithinWatermark(colNames) } @@ -2418,7 +2429,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 1.6.0 */ @scala.annotation.varargs - def describe(cols: String*): DS[Row] + def describe(cols: String*): Dataset[Row] /** * Computes specified statistics for numeric and string columns. Available statistics are:
    @@ -2488,7 +2499,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.3.0 */ @scala.annotation.varargs - def summary(statistics: String*): DS[Row] + def summary(statistics: String*): Dataset[Row] /** * Returns the first `n` rows. @@ -2520,7 +2531,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { /** * Concise syntax for chaining custom transformations. * {{{ - * def featurize(ds: DS[T]): DS[U] = ... + * def featurize(ds: Dataset[T]): Dataset[U] = ... * * ds * .transform(featurize) @@ -2530,7 +2541,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def transform[U](t: DS[T] => DS[U]): DS[U] = t(this.asInstanceOf[DS[T]]) + def transform[U](t: Dataset[T] => Dataset[U]): Dataset[U] = t(this.asInstanceOf[Dataset[T]]) /** * (Scala-specific) Returns a new Dataset that contains the result of applying `func` to each @@ -2539,7 +2550,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def map[U: Encoder](func: T => U): DS[U] + def map[U: Encoder](func: T => U): Dataset[U] /** * (Java-specific) Returns a new Dataset that contains the result of applying `func` to each @@ -2548,7 +2559,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def map[U](func: MapFunction[T, U], encoder: Encoder[U]): DS[U] + def map[U](func: MapFunction[T, U], encoder: Encoder[U]): Dataset[U] /** * (Scala-specific) Returns a new Dataset that contains the result of applying `func` to each @@ -2557,7 +2568,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def mapPartitions[U: Encoder](func: Iterator[T] => Iterator[U]): DS[U] + def mapPartitions[U: Encoder](func: Iterator[T] => Iterator[U]): Dataset[U] /** * (Java-specific) Returns a new Dataset that contains the result of applying `f` to each @@ -2566,7 +2577,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def mapPartitions[U](f: MapPartitionsFunction[T, U], encoder: Encoder[U]): DS[U] = + def mapPartitions[U](f: MapPartitionsFunction[T, U], encoder: Encoder[U]): Dataset[U] = mapPartitions(ToScalaUDF(f))(encoder) /** @@ -2576,7 +2587,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def flatMap[U: Encoder](func: T => IterableOnce[U]): DS[U] = + def flatMap[U: Encoder](func: T => IterableOnce[U]): Dataset[U] = mapPartitions(UDFAdaptors.flatMapToMapPartitions[T, U](func)) /** @@ -2586,7 +2597,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def flatMap[U](f: FlatMapFunction[T, U], encoder: Encoder[U]): DS[U] = { + def flatMap[U](f: FlatMapFunction[T, U], encoder: Encoder[U]): Dataset[U] = { mapPartitions(UDFAdaptors.flatMapToMapPartitions(f))(encoder) } @@ -2713,11 +2724,11 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def repartition(numPartitions: Int): DS[T] + def repartition(numPartitions: Int): Dataset[T] protected def repartitionByExpression( numPartitions: Option[Int], - partitionExprs: Seq[Column]): DS[T] + partitionExprs: Seq[Column]): Dataset[T] /** * Returns a new Dataset partitioned by the given partitioning expressions into `numPartitions`. @@ -2729,7 +2740,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def repartition(numPartitions: Int, partitionExprs: Column*): DS[T] = { + def repartition(numPartitions: Int, partitionExprs: Column*): Dataset[T] = { repartitionByExpression(Some(numPartitions), partitionExprs) } @@ -2744,11 +2755,13 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.0.0 */ @scala.annotation.varargs - def repartition(partitionExprs: Column*): DS[T] = { + def repartition(partitionExprs: Column*): Dataset[T] = { repartitionByExpression(None, partitionExprs) } - protected def repartitionByRange(numPartitions: Option[Int], partitionExprs: Seq[Column]): DS[T] + protected def repartitionByRange( + numPartitions: Option[Int], + partitionExprs: Seq[Column]): Dataset[T] /** * Returns a new Dataset partitioned by the given partitioning expressions into `numPartitions`. @@ -2766,7 +2779,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.3.0 */ @scala.annotation.varargs - def repartitionByRange(numPartitions: Int, partitionExprs: Column*): DS[T] = { + def repartitionByRange(numPartitions: Int, partitionExprs: Column*): Dataset[T] = { repartitionByRange(Some(numPartitions), partitionExprs) } @@ -2787,7 +2800,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @since 2.3.0 */ @scala.annotation.varargs - def repartitionByRange(partitionExprs: Column*): DS[T] = { + def repartitionByRange(partitionExprs: Column*): Dataset[T] = { repartitionByRange(None, partitionExprs) } @@ -2807,7 +2820,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 1.6.0 */ - def coalesce(numPartitions: Int): DS[T] + def coalesce(numPartitions: Int): Dataset[T] /** * Returns a new Dataset that contains only the unique rows from this Dataset. This is an alias @@ -2823,7 +2836,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group typedrel * @since 2.0.0 */ - def distinct(): DS[T] = dropDuplicates() + def distinct(): Dataset[T] = dropDuplicates() /** * Persist this Dataset with the default storage level (`MEMORY_AND_DISK`). @@ -2831,7 +2844,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def persist(): DS[T] + def persist(): Dataset[T] /** * Persist this Dataset with the default storage level (`MEMORY_AND_DISK`). @@ -2839,7 +2852,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def cache(): DS[T] + def cache(): Dataset[T] /** * Persist this Dataset with the given storage level. @@ -2850,7 +2863,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def persist(newLevel: StorageLevel): DS[T] + def persist(newLevel: StorageLevel): Dataset[T] /** * Get the Dataset's current storage level, or StorageLevel.NONE if not persisted. @@ -2869,7 +2882,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def unpersist(blocking: Boolean): DS[T] + def unpersist(blocking: Boolean): Dataset[T] /** * Mark the Dataset as non-persistent, and remove all blocks for it from memory and disk. This @@ -2878,7 +2891,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * @group basic * @since 1.6.0 */ - def unpersist(): DS[T] + def unpersist(): Dataset[T] /** * Registers this Dataset as a temporary table using the given name. The lifetime of this @@ -3008,7 +3021,7 @@ abstract class Dataset[T, DS[U] <: Dataset[U, DS]] extends Serializable { * * @since 2.0.0 */ - def toJSON: DS[String] + def toJSON: Dataset[String] /** * Returns a best-effort snapshot of the files that compose this Dataset. This method simply diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/KeyValueGroupedDataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/KeyValueGroupedDataset.scala index 50dfbff81dd3e..81f999430a128 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/KeyValueGroupedDataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/KeyValueGroupedDataset.scala @@ -30,8 +30,8 @@ import org.apache.spark.sql.streaming.{GroupState, GroupStateTimeout, OutputMode * * @since 2.0.0 */ -abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Serializable { - type KVDS[KY, VL] <: KeyValueGroupedDataset[KY, VL, DS] +abstract class KeyValueGroupedDataset[K, V] extends Serializable { + type KVDS[KL, VL] <: KeyValueGroupedDataset[KL, VL] /** * Returns a new [[KeyValueGroupedDataset]] where the type of the key has been mapped to the @@ -40,7 +40,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def keyAs[L: Encoder]: KVDS[L, V] + def keyAs[L: Encoder]: KeyValueGroupedDataset[L, V] /** * Returns a new [[KeyValueGroupedDataset]] where the given function `func` has been applied to @@ -53,7 +53,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 2.1.0 */ - def mapValues[W: Encoder](func: V => W): KVDS[K, W] + def mapValues[W: Encoder](func: V => W): KeyValueGroupedDataset[K, W] /** * Returns a new [[KeyValueGroupedDataset]] where the given function `func` has been applied to @@ -68,7 +68,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 2.1.0 */ - def mapValues[W](func: MapFunction[V, W], encoder: Encoder[W]): KVDS[K, W] = { + def mapValues[W](func: MapFunction[V, W], encoder: Encoder[W]): KeyValueGroupedDataset[K, W] = { mapValues(ToScalaUDF(func))(encoder) } @@ -78,7 +78,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def keys: DS[K] + def keys: Dataset[K] /** * (Scala-specific) Applies the given function to each group of data. For each unique group, the @@ -98,7 +98,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def flatMapGroups[U: Encoder](f: (K, Iterator[V]) => IterableOnce[U]): DS[U] = { + def flatMapGroups[U: Encoder](f: (K, Iterator[V]) => IterableOnce[U]): Dataset[U] = { flatMapSortedGroups(Nil: _*)(f) } @@ -120,7 +120,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def flatMapGroups[U](f: FlatMapGroupsFunction[K, V, U], encoder: Encoder[U]): DS[U] = { + def flatMapGroups[U](f: FlatMapGroupsFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = { flatMapGroups(ToScalaUDF(f))(encoder) } @@ -149,7 +149,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * @since 3.4.0 */ def flatMapSortedGroups[U: Encoder](sortExprs: Column*)( - f: (K, Iterator[V]) => IterableOnce[U]): DS[U] + f: (K, Iterator[V]) => IterableOnce[U]): Dataset[U] /** * (Java-specific) Applies the given function to each group of data. For each unique group, the @@ -178,7 +178,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser def flatMapSortedGroups[U]( SortExprs: Array[Column], f: FlatMapGroupsFunction[K, V, U], - encoder: Encoder[U]): DS[U] = { + encoder: Encoder[U]): Dataset[U] = { import org.apache.spark.util.ArrayImplicits._ flatMapSortedGroups(SortExprs.toImmutableArraySeq: _*)(ToScalaUDF(f))(encoder) } @@ -201,7 +201,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def mapGroups[U: Encoder](f: (K, Iterator[V]) => U): DS[U] = { + def mapGroups[U: Encoder](f: (K, Iterator[V]) => U): Dataset[U] = { flatMapGroups(UDFAdaptors.mapGroupsToFlatMapGroups(f)) } @@ -223,7 +223,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def mapGroups[U](f: MapGroupsFunction[K, V, U], encoder: Encoder[U]): DS[U] = { + def mapGroups[U](f: MapGroupsFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = { mapGroups(ToScalaUDF(f))(encoder) } @@ -247,7 +247,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * @since 2.2.0 */ def mapGroupsWithState[S: Encoder, U: Encoder]( - func: (K, Iterator[V], GroupState[S]) => U): DS[U] + func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] /** * (Scala-specific) Applies the given function to each group of data, while maintaining a @@ -271,7 +271,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * @since 2.2.0 */ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)( - func: (K, Iterator[V], GroupState[S]) => U): DS[U] + func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] /** * (Scala-specific) Applies the given function to each group of data, while maintaining a @@ -301,7 +301,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser */ def mapGroupsWithState[S: Encoder, U: Encoder]( timeoutConf: GroupStateTimeout, - initialState: KVDS[K, S])(func: (K, Iterator[V], GroupState[S]) => U): DS[U] + initialState: KVDS[K, S])(func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] /** * (Java-specific) Applies the given function to each group of data, while maintaining a @@ -329,7 +329,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser def mapGroupsWithState[S, U]( func: MapGroupsWithStateFunction[K, V, S, U], stateEncoder: Encoder[S], - outputEncoder: Encoder[U]): DS[U] = { + outputEncoder: Encoder[U]): Dataset[U] = { mapGroupsWithState[S, U](ToScalaUDF(func))(stateEncoder, outputEncoder) } @@ -362,7 +362,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser func: MapGroupsWithStateFunction[K, V, S, U], stateEncoder: Encoder[S], outputEncoder: Encoder[U], - timeoutConf: GroupStateTimeout): DS[U] = { + timeoutConf: GroupStateTimeout): Dataset[U] = { mapGroupsWithState[S, U](timeoutConf)(ToScalaUDF(func))(stateEncoder, outputEncoder) } @@ -400,7 +400,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser stateEncoder: Encoder[S], outputEncoder: Encoder[U], timeoutConf: GroupStateTimeout, - initialState: KVDS[K, S]): DS[U] = { + initialState: KVDS[K, S]): Dataset[U] = { val f = ToScalaUDF(func) mapGroupsWithState[S, U](timeoutConf, initialState)(f)(stateEncoder, outputEncoder) } @@ -430,7 +430,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser */ def flatMapGroupsWithState[S: Encoder, U: Encoder]( outputMode: OutputMode, - timeoutConf: GroupStateTimeout)(func: (K, Iterator[V], GroupState[S]) => Iterator[U]): DS[U] + timeoutConf: GroupStateTimeout)( + func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] /** * (Scala-specific) Applies the given function to each group of data, while maintaining a @@ -462,7 +463,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser def flatMapGroupsWithState[S: Encoder, U: Encoder]( outputMode: OutputMode, timeoutConf: GroupStateTimeout, - initialState: KVDS[K, S])(func: (K, Iterator[V], GroupState[S]) => Iterator[U]): DS[U] + initialState: KVDS[K, S])(func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] /** * (Java-specific) Applies the given function to each group of data, while maintaining a @@ -496,7 +497,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser outputMode: OutputMode, stateEncoder: Encoder[S], outputEncoder: Encoder[U], - timeoutConf: GroupStateTimeout): DS[U] = { + timeoutConf: GroupStateTimeout): Dataset[U] = { val f = ToScalaUDF(func) flatMapGroupsWithState[S, U](outputMode, timeoutConf)(f)(stateEncoder, outputEncoder) } @@ -540,7 +541,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser stateEncoder: Encoder[S], outputEncoder: Encoder[U], timeoutConf: GroupStateTimeout, - initialState: KVDS[K, S]): DS[U] = { + initialState: KVDS[K, S]): Dataset[U] = { flatMapGroupsWithState[S, U](outputMode, timeoutConf, initialState)(ToScalaUDF(func))( stateEncoder, outputEncoder) @@ -568,7 +569,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser private[sql] def transformWithState[U: Encoder]( statefulProcessor: StatefulProcessor[K, V, U], timeMode: TimeMode, - outputMode: OutputMode): DS[U] + outputMode: OutputMode): Dataset[U] /** * (Scala-specific) Invokes methods defined in the stateful processor used in arbitrary state @@ -597,7 +598,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser private[sql] def transformWithState[U: Encoder]( statefulProcessor: StatefulProcessor[K, V, U], eventTimeColumnName: String, - outputMode: OutputMode): DS[U] + outputMode: OutputMode): Dataset[U] /** * (Java-specific) Invokes methods defined in the stateful processor used in arbitrary state API @@ -624,7 +625,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser statefulProcessor: StatefulProcessor[K, V, U], timeMode: TimeMode, outputMode: OutputMode, - outputEncoder: Encoder[U]): DS[U] = { + outputEncoder: Encoder[U]): Dataset[U] = { transformWithState(statefulProcessor, timeMode, outputMode)(outputEncoder) } @@ -660,7 +661,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser statefulProcessor: StatefulProcessor[K, V, U], eventTimeColumnName: String, outputMode: OutputMode, - outputEncoder: Encoder[U]): DS[U] = { + outputEncoder: Encoder[U]): Dataset[U] = { transformWithState(statefulProcessor, eventTimeColumnName, outputMode)(outputEncoder) } @@ -689,7 +690,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser statefulProcessor: StatefulProcessorWithInitialState[K, V, U, S], timeMode: TimeMode, outputMode: OutputMode, - initialState: KVDS[K, S]): DS[U] + initialState: KVDS[K, S]): Dataset[U] /** * (Scala-specific) Invokes methods defined in the stateful processor used in arbitrary state @@ -722,7 +723,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser statefulProcessor: StatefulProcessorWithInitialState[K, V, U, S], eventTimeColumnName: String, outputMode: OutputMode, - initialState: KVDS[K, S]): DS[U] + initialState: KVDS[K, S]): Dataset[U] /** * (Java-specific) Invokes methods defined in the stateful processor used in arbitrary state API @@ -756,7 +757,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser outputMode: OutputMode, initialState: KVDS[K, S], outputEncoder: Encoder[U], - initialStateEncoder: Encoder[S]): DS[U] = { + initialStateEncoder: Encoder[S]): Dataset[U] = { transformWithState(statefulProcessor, timeMode, outputMode, initialState)( outputEncoder, initialStateEncoder) @@ -798,7 +799,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser initialState: KVDS[K, S], eventTimeColumnName: String, outputEncoder: Encoder[U], - initialStateEncoder: Encoder[S]): DS[U] = { + initialStateEncoder: Encoder[S]): Dataset[U] = { transformWithState(statefulProcessor, eventTimeColumnName, outputMode, initialState)( outputEncoder, initialStateEncoder) @@ -811,7 +812,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def reduceGroups(f: (V, V) => V): DS[(K, V)] + def reduceGroups(f: (V, V) => V): Dataset[(K, V)] /** * (Java-specific) Reduces the elements of each group of data using the specified binary @@ -820,7 +821,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def reduceGroups(f: ReduceFunction[V]): DS[(K, V)] = { + def reduceGroups(f: ReduceFunction[V]): Dataset[(K, V)] = { reduceGroups(ToScalaUDF(f)) } @@ -829,7 +830,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * and code reuse, we do this without the help of the type system and then use helper functions * that cast appropriately for the user facing interface. */ - protected def aggUntyped(columns: TypedColumn[_, _]*): DS[_] + protected def aggUntyped(columns: TypedColumn[_, _]*): Dataset[_] /** * Computes the given aggregation, returning a [[Dataset]] of tuples for each unique key and the @@ -837,8 +838,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def agg[U1](col1: TypedColumn[V, U1]): DS[(K, U1)] = - aggUntyped(col1).asInstanceOf[DS[(K, U1)]] + def agg[U1](col1: TypedColumn[V, U1]): Dataset[(K, U1)] = + aggUntyped(col1).asInstanceOf[Dataset[(K, U1)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -846,8 +847,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def agg[U1, U2](col1: TypedColumn[V, U1], col2: TypedColumn[V, U2]): DS[(K, U1, U2)] = - aggUntyped(col1, col2).asInstanceOf[DS[(K, U1, U2)]] + def agg[U1, U2](col1: TypedColumn[V, U1], col2: TypedColumn[V, U2]): Dataset[(K, U1, U2)] = + aggUntyped(col1, col2).asInstanceOf[Dataset[(K, U1, U2)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -858,8 +859,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser def agg[U1, U2, U3]( col1: TypedColumn[V, U1], col2: TypedColumn[V, U2], - col3: TypedColumn[V, U3]): DS[(K, U1, U2, U3)] = - aggUntyped(col1, col2, col3).asInstanceOf[DS[(K, U1, U2, U3)]] + col3: TypedColumn[V, U3]): Dataset[(K, U1, U2, U3)] = + aggUntyped(col1, col2, col3).asInstanceOf[Dataset[(K, U1, U2, U3)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -871,8 +872,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser col1: TypedColumn[V, U1], col2: TypedColumn[V, U2], col3: TypedColumn[V, U3], - col4: TypedColumn[V, U4]): DS[(K, U1, U2, U3, U4)] = - aggUntyped(col1, col2, col3, col4).asInstanceOf[DS[(K, U1, U2, U3, U4)]] + col4: TypedColumn[V, U4]): Dataset[(K, U1, U2, U3, U4)] = + aggUntyped(col1, col2, col3, col4).asInstanceOf[Dataset[(K, U1, U2, U3, U4)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -885,8 +886,8 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser col2: TypedColumn[V, U2], col3: TypedColumn[V, U3], col4: TypedColumn[V, U4], - col5: TypedColumn[V, U5]): DS[(K, U1, U2, U3, U4, U5)] = - aggUntyped(col1, col2, col3, col4, col5).asInstanceOf[DS[(K, U1, U2, U3, U4, U5)]] + col5: TypedColumn[V, U5]): Dataset[(K, U1, U2, U3, U4, U5)] = + aggUntyped(col1, col2, col3, col4, col5).asInstanceOf[Dataset[(K, U1, U2, U3, U4, U5)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -900,9 +901,9 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser col3: TypedColumn[V, U3], col4: TypedColumn[V, U4], col5: TypedColumn[V, U5], - col6: TypedColumn[V, U6]): DS[(K, U1, U2, U3, U4, U5, U6)] = + col6: TypedColumn[V, U6]): Dataset[(K, U1, U2, U3, U4, U5, U6)] = aggUntyped(col1, col2, col3, col4, col5, col6) - .asInstanceOf[DS[(K, U1, U2, U3, U4, U5, U6)]] + .asInstanceOf[Dataset[(K, U1, U2, U3, U4, U5, U6)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -917,9 +918,9 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser col4: TypedColumn[V, U4], col5: TypedColumn[V, U5], col6: TypedColumn[V, U6], - col7: TypedColumn[V, U7]): DS[(K, U1, U2, U3, U4, U5, U6, U7)] = + col7: TypedColumn[V, U7]): Dataset[(K, U1, U2, U3, U4, U5, U6, U7)] = aggUntyped(col1, col2, col3, col4, col5, col6, col7) - .asInstanceOf[DS[(K, U1, U2, U3, U4, U5, U6, U7)]] + .asInstanceOf[Dataset[(K, U1, U2, U3, U4, U5, U6, U7)]] /** * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key and @@ -935,9 +936,9 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser col5: TypedColumn[V, U5], col6: TypedColumn[V, U6], col7: TypedColumn[V, U7], - col8: TypedColumn[V, U8]): DS[(K, U1, U2, U3, U4, U5, U6, U7, U8)] = + col8: TypedColumn[V, U8]): Dataset[(K, U1, U2, U3, U4, U5, U6, U7, U8)] = aggUntyped(col1, col2, col3, col4, col5, col6, col7, col8) - .asInstanceOf[DS[(K, U1, U2, U3, U4, U5, U6, U7, U8)]] + .asInstanceOf[Dataset[(K, U1, U2, U3, U4, U5, U6, U7, U8)]] /** * Returns a [[Dataset]] that contains a tuple with each key and the number of items present for @@ -945,7 +946,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * * @since 1.6.0 */ - def count(): DS[(K, Long)] = agg(cnt(lit(1)).as(PrimitiveLongEncoder)) + def count(): Dataset[(K, Long)] = agg(cnt(lit(1)).as(PrimitiveLongEncoder)) /** * (Scala-specific) Applies the given function to each cogrouped data. For each unique group, @@ -956,7 +957,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * @since 1.6.0 */ def cogroup[U, R: Encoder](other: KVDS[K, U])( - f: (K, Iterator[V], Iterator[U]) => IterableOnce[R]): DS[R] = { + f: (K, Iterator[V], Iterator[U]) => IterableOnce[R]): Dataset[R] = { cogroupSorted(other)(Nil: _*)(Nil: _*)(f) } @@ -971,7 +972,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser def cogroup[U, R]( other: KVDS[K, U], f: CoGroupFunction[K, V, U, R], - encoder: Encoder[R]): DS[R] = { + encoder: Encoder[R]): Dataset[R] = { cogroup(other)(ToScalaUDF(f))(encoder) } @@ -991,7 +992,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser * @since 3.4.0 */ def cogroupSorted[U, R: Encoder](other: KVDS[K, U])(thisSortExprs: Column*)( - otherSortExprs: Column*)(f: (K, Iterator[V], Iterator[U]) => IterableOnce[R]): DS[R] + otherSortExprs: Column*)(f: (K, Iterator[V], Iterator[U]) => IterableOnce[R]): Dataset[R] /** * (Java-specific) Applies the given function to each sorted cogrouped data. For each unique @@ -1013,7 +1014,7 @@ abstract class KeyValueGroupedDataset[K, V, DS[U] <: Dataset[U, DS]] extends Ser thisSortExprs: Array[Column], otherSortExprs: Array[Column], f: CoGroupFunction[K, V, U, R], - encoder: Encoder[R]): DS[R] = { + encoder: Encoder[R]): Dataset[R] = { import org.apache.spark.util.ArrayImplicits._ cogroupSorted(other)(thisSortExprs.toImmutableArraySeq: _*)( otherSortExprs.toImmutableArraySeq: _*)(ToScalaUDF(f))(encoder) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/RelationalGroupedDataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/RelationalGroupedDataset.scala index 7dd5f46beb316..118b8f1ecd488 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/RelationalGroupedDataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/RelationalGroupedDataset.scala @@ -35,15 +35,13 @@ import org.apache.spark.sql.{functions, Column, Encoder, Row} * @since 2.0.0 */ @Stable -abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { - type RGD <: RelationalGroupedDataset[DS] - - protected def df: DS[Row] +abstract class RelationalGroupedDataset { + protected def df: Dataset[Row] /** * Create a aggregation based on the grouping column, the grouping type, and the aggregations. */ - protected def toDF(aggCols: Seq[Column]): DS[Row] + protected def toDF(aggCols: Seq[Column]): Dataset[Row] protected def selectNumericColumns(colNames: Seq[String]): Seq[Column] @@ -62,7 +60,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { private def aggregateNumericColumns( colNames: Seq[String], - function: Column => Column): DS[Row] = { + function: Column => Column): Dataset[Row] = { toDF(selectNumericColumns(colNames).map(function)) } @@ -72,7 +70,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * * @since 3.0.0 */ - def as[K: Encoder, T: Encoder]: KeyValueGroupedDataset[K, T, DS] + def as[K: Encoder, T: Encoder]: KeyValueGroupedDataset[K, T] /** * (Scala-specific) Compute aggregates by specifying the column names and aggregate methods. The @@ -89,7 +87,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * * @since 1.3.0 */ - def agg(aggExpr: (String, String), aggExprs: (String, String)*): DS[Row] = + def agg(aggExpr: (String, String), aggExprs: (String, String)*): Dataset[Row] = toDF((aggExpr +: aggExprs).map(toAggCol)) /** @@ -107,7 +105,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * * @since 1.3.0 */ - def agg(exprs: Map[String, String]): DS[Row] = toDF(exprs.map(toAggCol).toSeq) + def agg(exprs: Map[String, String]): Dataset[Row] = toDF(exprs.map(toAggCol).toSeq) /** * (Java-specific) Compute aggregates by specifying a map from column name to aggregate methods. @@ -122,7 +120,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * * @since 1.3.0 */ - def agg(exprs: util.Map[String, String]): DS[Row] = { + def agg(exprs: util.Map[String, String]): Dataset[Row] = { agg(exprs.asScala.toMap) } @@ -158,7 +156,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def agg(expr: Column, exprs: Column*): DS[Row] = toDF(expr +: exprs) + def agg(expr: Column, exprs: Column*): Dataset[Row] = toDF(expr +: exprs) /** * Count the number of rows for each group. The resulting `DataFrame` will also contain the @@ -166,7 +164,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * * @since 1.3.0 */ - def count(): DS[Row] = toDF(functions.count(functions.lit(1)).as("count") :: Nil) + def count(): Dataset[Row] = toDF(functions.count(functions.lit(1)).as("count") :: Nil) /** * Compute the average value for each numeric columns for each group. This is an alias for @@ -176,7 +174,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def mean(colNames: String*): DS[Row] = aggregateNumericColumns(colNames, functions.avg) + def mean(colNames: String*): Dataset[Row] = aggregateNumericColumns(colNames, functions.avg) /** * Compute the max value for each numeric columns for each group. The resulting `DataFrame` will @@ -186,7 +184,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def max(colNames: String*): DS[Row] = aggregateNumericColumns(colNames, functions.max) + def max(colNames: String*): Dataset[Row] = aggregateNumericColumns(colNames, functions.max) /** * Compute the mean value for each numeric columns for each group. The resulting `DataFrame` @@ -196,7 +194,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def avg(colNames: String*): DS[Row] = aggregateNumericColumns(colNames, functions.avg) + def avg(colNames: String*): Dataset[Row] = aggregateNumericColumns(colNames, functions.avg) /** * Compute the min value for each numeric column for each group. The resulting `DataFrame` will @@ -206,7 +204,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def min(colNames: String*): DS[Row] = aggregateNumericColumns(colNames, functions.min) + def min(colNames: String*): Dataset[Row] = aggregateNumericColumns(colNames, functions.min) /** * Compute the sum for each numeric columns for each group. The resulting `DataFrame` will also @@ -216,7 +214,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * @since 1.3.0 */ @scala.annotation.varargs - def sum(colNames: String*): DS[Row] = aggregateNumericColumns(colNames, functions.sum) + def sum(colNames: String*): Dataset[Row] = aggregateNumericColumns(colNames, functions.sum) /** * Pivots a column of the current `DataFrame` and performs the specified aggregation. @@ -237,7 +235,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * Name of the column to pivot. * @since 1.6.0 */ - def pivot(pivotColumn: String): RGD = pivot(df.col(pivotColumn)) + def pivot(pivotColumn: String): RelationalGroupedDataset = pivot(df.col(pivotColumn)) /** * Pivots a column of the current `DataFrame` and performs the specified aggregation. There are @@ -271,7 +269,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * List of values that will be translated to columns in the output DataFrame. * @since 1.6.0 */ - def pivot(pivotColumn: String, values: Seq[Any]): RGD = + def pivot(pivotColumn: String, values: Seq[Any]): RelationalGroupedDataset = pivot(df.col(pivotColumn), values) /** @@ -299,7 +297,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * List of values that will be translated to columns in the output DataFrame. * @since 1.6.0 */ - def pivot(pivotColumn: String, values: util.List[Any]): RGD = + def pivot(pivotColumn: String, values: util.List[Any]): RelationalGroupedDataset = pivot(df.col(pivotColumn), values) /** @@ -316,7 +314,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * List of values that will be translated to columns in the output DataFrame. * @since 2.4.0 */ - def pivot(pivotColumn: Column, values: util.List[Any]): RGD = + def pivot(pivotColumn: Column, values: util.List[Any]): RelationalGroupedDataset = pivot(pivotColumn, values.asScala.toSeq) /** @@ -338,7 +336,7 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * he column to pivot. * @since 2.4.0 */ - def pivot(pivotColumn: Column): RGD + def pivot(pivotColumn: Column): RelationalGroupedDataset /** * Pivots a column of the current `DataFrame` and performs the specified aggregation. This is an @@ -358,5 +356,5 @@ abstract class RelationalGroupedDataset[DS[U] <: Dataset[U, DS]] { * List of values that will be translated to columns in the output DataFrame. * @since 2.4.0 */ - def pivot(pivotColumn: Column, values: Seq[Any]): RGD + def pivot(pivotColumn: Column, values: Seq[Any]): RelationalGroupedDataset } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 63d4a12e11839..41d16b16ab1c5 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -49,7 +49,7 @@ import org.apache.spark.sql.types.StructType * .getOrCreate() * }}} */ -abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with Closeable { +abstract class SparkSession extends Serializable with Closeable { /** * The version of Spark on which this application is running. @@ -103,7 +103,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * will initialize the metastore, which may take some time. * @since 2.0.0 */ - def newSession(): SparkSession[DS] + def newSession(): SparkSession /* --------------------------------- * | Methods for creating DataFrames | @@ -115,14 +115,14 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * @since 2.0.0 */ @transient - def emptyDataFrame: DS[Row] + def emptyDataFrame: Dataset[Row] /** * Creates a `DataFrame` from a local Seq of Product. * * @since 2.0.0 */ - def createDataFrame[A <: Product: TypeTag](data: Seq[A]): DS[Row] + def createDataFrame[A <: Product: TypeTag](data: Seq[A]): Dataset[Row] /** * :: DeveloperApi :: Creates a `DataFrame` from a `java.util.List` containing @@ -133,7 +133,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * @since 2.0.0 */ @DeveloperApi - def createDataFrame(rows: util.List[Row], schema: StructType): DS[Row] + def createDataFrame(rows: util.List[Row], schema: StructType): Dataset[Row] /** * Applies a schema to a List of Java Beans. @@ -143,7 +143,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 1.6.0 */ - def createDataFrame(data: util.List[_], beanClass: Class[_]): DS[Row] + def createDataFrame(data: util.List[_], beanClass: Class[_]): Dataset[Row] /* ------------------------------- * | Methods for creating DataSets | @@ -154,7 +154,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def emptyDataset[T: Encoder]: DS[T] + def emptyDataset[T: Encoder]: Dataset[T] /** * Creates a [[Dataset]] from a local Seq of data of a given type. This method requires an @@ -183,7 +183,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def createDataset[T: Encoder](data: Seq[T]): DS[T] + def createDataset[T: Encoder](data: Seq[T]): Dataset[T] /** * Creates a [[Dataset]] from a `java.util.List` of a given type. This method requires an @@ -200,7 +200,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def createDataset[T: Encoder](data: util.List[T]): DS[T] + def createDataset[T: Encoder](data: util.List[T]): Dataset[T] /** * Creates a [[Dataset]] with a single `LongType` column named `id`, containing elements in a @@ -208,7 +208,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def range(end: Long): DS[lang.Long] + def range(end: Long): Dataset[lang.Long] /** * Creates a [[Dataset]] with a single `LongType` column named `id`, containing elements in a @@ -216,7 +216,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def range(start: Long, end: Long): DS[lang.Long] + def range(start: Long, end: Long): Dataset[lang.Long] /** * Creates a [[Dataset]] with a single `LongType` column named `id`, containing elements in a @@ -224,7 +224,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def range(start: Long, end: Long, step: Long): DS[lang.Long] + def range(start: Long, end: Long, step: Long): Dataset[lang.Long] /** * Creates a [[Dataset]] with a single `LongType` column named `id`, containing elements in a @@ -232,7 +232,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def range(start: Long, end: Long, step: Long, numPartitions: Int): DS[lang.Long] + def range(start: Long, end: Long, step: Long, numPartitions: Int): Dataset[lang.Long] /* ------------------------- * | Catalog-related methods | @@ -244,7 +244,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def catalog: Catalog[DS] + def catalog: Catalog /** * Returns the specified table/view as a `DataFrame`. If it's a table, it must support batch @@ -259,7 +259,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * database. Note that, the global temporary view database is also valid here. * @since 2.0.0 */ - def table(tableName: String): DS[Row] + def table(tableName: String): Dataset[Row] /* ----------------- * | Everything else | @@ -281,7 +281,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * @since 3.5.0 */ @Experimental - def sql(sqlText: String, args: Array[_]): DS[Row] + def sql(sqlText: String, args: Array[_]): Dataset[Row] /** * Executes a SQL query substituting named parameters by the given arguments, returning the @@ -299,7 +299,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * @since 3.4.0 */ @Experimental - def sql(sqlText: String, args: Map[String, Any]): DS[Row] + def sql(sqlText: String, args: Map[String, Any]): Dataset[Row] /** * Executes a SQL query substituting named parameters by the given arguments, returning the @@ -317,7 +317,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * @since 3.4.0 */ @Experimental - def sql(sqlText: String, args: util.Map[String, Any]): DS[Row] = { + def sql(sqlText: String, args: util.Map[String, Any]): Dataset[Row] = { sql(sqlText, args.asScala.toMap) } @@ -327,7 +327,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def sql(sqlText: String): DS[Row] = sql(sqlText, Map.empty[String, Any]) + def sql(sqlText: String): Dataset[Row] = sql(sqlText, Map.empty[String, Any]) /** * Add a single artifact to the current session. @@ -503,7 +503,7 @@ abstract class SparkSession[DS[U] <: Dataset[U, DS]] extends Serializable with C * * @since 2.0.0 */ - def read: DataFrameReader[DS] + def read: DataFrameReader /** * Executes some code block and prints to stdout the time taken to execute the block. This is diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQuery.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQuery.scala index 16cd45339f051..0aeb3518facd8 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQuery.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQuery.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.streaming.{StreamingQueryException, StreamingQueryPr * @since 2.0.0 */ @Evolving -trait StreamingQuery[DS[U] <: Dataset[U, DS]] { +trait StreamingQuery { /** * Returns the user-specified name of the query, or null if not specified. This name can be @@ -62,7 +62,7 @@ trait StreamingQuery[DS[U] <: Dataset[U, DS]] { * * @since 2.0.0 */ - def sparkSession: SparkSession[DS] + def sparkSession: SparkSession /** * Returns `true` if this query is actively running. diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 86f8923f36b40..02669270c8acf 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -1714,7 +1714,7 @@ object functions { * @group normal_funcs * @since 1.5.0 */ - def broadcast[DS[U] <: api.Dataset[U, DS]](df: DS[_]): df.type = { + def broadcast[DS[U] <: api.Dataset[U]](df: DS[_]): df.type = { df.hint("broadcast").asInstanceOf[df.type] } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala index 53640f513fc81..b356751083fc1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala @@ -21,6 +21,7 @@ import java.{lang => jl} import org.apache.spark.annotation.Stable import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.ExpressionUtils.column @@ -33,7 +34,7 @@ import org.apache.spark.sql.types._ */ @Stable final class DataFrameNaFunctions private[sql](df: DataFrame) - extends api.DataFrameNaFunctions[Dataset] { + extends api.DataFrameNaFunctions { import df.sparkSession.RichColumn protected def drop(minNonNulls: Option[Int]): Dataset[Row] = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala index f105a77cf253b..78cc65bb7a298 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.ExprUtils import org.apache.spark.sql.catalyst.json.{CreateJacksonParser, JacksonParser, JSONOptions} import org.apache.spark.sql.catalyst.util.FailureSafeParser import org.apache.spark.sql.catalyst.xml.{StaxXmlParser, XmlOptions} +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.DDLUtils import org.apache.spark.sql.execution.datasources.DataSource @@ -54,7 +55,9 @@ import org.apache.spark.unsafe.types.UTF8String */ @Stable class DataFrameReader private[sql](sparkSession: SparkSession) - extends api.DataFrameReader[Dataset] { + extends api.DataFrameReader { + override type DS[U] = Dataset[U] + format(sparkSession.sessionState.conf.defaultDataSourceName) /** @inheritdoc */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala index a5ab237bb7041..9f7180d8dfd6a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala @@ -23,6 +23,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Stable import org.apache.spark.sql.catalyst.trees.CurrentOrigin.withOrigin +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.execution.stat._ import org.apache.spark.sql.functions.col import org.apache.spark.util.ArrayImplicits._ @@ -34,7 +35,7 @@ import org.apache.spark.util.ArrayImplicits._ */ @Stable final class DataFrameStatFunctions private[sql](protected val df: DataFrame) - extends api.DataFrameStatFunctions[Dataset] { + extends api.DataFrameStatFunctions { /** @inheritdoc */ def approxQuantile( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index c147b6a56e024..61f9e6ff7c042 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -52,6 +52,7 @@ import org.apache.spark.sql.catalyst.trees.{TreeNodeTag, TreePattern} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, IntervalUtils} import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression @@ -215,7 +216,8 @@ private[sql] object Dataset { class Dataset[T] private[sql]( @DeveloperApi @Unstable @transient val queryExecution: QueryExecution, @DeveloperApi @Unstable @transient val encoder: Encoder[T]) - extends api.Dataset[T, Dataset] { + extends api.Dataset[T] { + type DS[U] = Dataset[U] type RGD = RelationalGroupedDataset @transient lazy val sparkSession: SparkSession = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala index fcad1b721eaca..c645ba57e8f82 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderF import org.apache.spark.sql.catalyst.encoders.encoderFor import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.expressions.ReduceAggregator import org.apache.spark.sql.internal.TypedAggUtils.{aggKeyColumn, withInputType} @@ -41,7 +42,7 @@ class KeyValueGroupedDataset[K, V] private[sql]( @transient val queryExecution: QueryExecution, private val dataAttributes: Seq[Attribute], private val groupingAttributes: Seq[Attribute]) - extends api.KeyValueGroupedDataset[K, V, Dataset] { + extends api.KeyValueGroupedDataset[K, V] { type KVDS[KY, VL] = KeyValueGroupedDataset[KY, VL] private implicit def kEncoderImpl: Encoder[K] = kEncoder diff --git a/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala index da4609135fd63..bd47a21a1e09b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.streaming.InternalOutputModes import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.toPrettySQL +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.internal.ExpressionUtils.{column, generateAlias} @@ -52,8 +53,8 @@ class RelationalGroupedDataset protected[sql]( protected[sql] val df: DataFrame, private[sql] val groupingExprs: Seq[Expression], groupType: RelationalGroupedDataset.GroupType) - extends api.RelationalGroupedDataset[Dataset] { - type RGD = RelationalGroupedDataset + extends api.RelationalGroupedDataset { + import RelationalGroupedDataset._ import df.sparkSession._ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index 720b77b0b9fe5..137dbaed9f00a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -96,7 +96,7 @@ class SparkSession private( @transient private[sql] val extensions: SparkSessionExtensions, @transient private[sql] val initialSessionOptions: Map[String, String], @transient private val parentManagedJobTags: Map[String, String]) - extends api.SparkSession[Dataset] with Logging { self => + extends api.SparkSession with Logging { self => // The call site where this SparkSession was constructed. private val creationSite: CallSite = Utils.getCallSite() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala index 661e43fe73cae..c39018ff06fca 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala @@ -19,10 +19,11 @@ package org.apache.spark.sql.catalog import java.util import org.apache.spark.sql.{api, DataFrame, Dataset} +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.types.StructType /** @inheritdoc */ -abstract class Catalog extends api.Catalog[Dataset] { +abstract class Catalog extends api.Catalog { /** @inheritdoc */ override def listDatabases(): Dataset[Database] diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala new file mode 100644 index 0000000000000..af91b57a6848b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.classic + +import scala.language.implicitConversions + +import org.apache.spark.annotation.DeveloperApi +import org.apache.spark.sql._ + +/** + * Conversions from sql interfaces to the Classic specific implementation. + * + * This class is mainly used by the implementation, but is also meant to be used by extension + * developers. + * + * We provide both a trait and an object. The trait is useful in situations where an extension + * developer needs to use these conversions in a project covering multiple Spark versions. They can + * create a shim for these conversions, the Spark 4+ version of the shim implements this trait, and + * shims for older versions do not. + */ +@DeveloperApi +trait ClassicConversions { + implicit def castToImpl(session: api.SparkSession): SparkSession = + session.asInstanceOf[SparkSession] + + implicit def castToImpl[T](ds: api.Dataset[T]): Dataset[T] = + ds.asInstanceOf[Dataset[T]] + + implicit def castToImpl(rgds: api.RelationalGroupedDataset): RelationalGroupedDataset = + rgds.asInstanceOf[RelationalGroupedDataset] + + implicit def castToImpl[K, V](kvds: api.KeyValueGroupedDataset[K, V]) + : KeyValueGroupedDataset[K, V] = kvds.asInstanceOf[KeyValueGroupedDataset[K, V]] +} + +object ClassicConversions extends ClassicConversions diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala index 653e1df4af679..7cf92db59067c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala @@ -16,10 +16,10 @@ */ package org.apache.spark.sql.streaming -import org.apache.spark.sql.{api, Dataset, SparkSession} +import org.apache.spark.sql.{api, SparkSession} /** @inheritdoc */ -trait StreamingQuery extends api.StreamingQuery[Dataset] { +trait StreamingQuery extends api.StreamingQuery { /** @inheritdoc */ override def sparkSession: SparkSession } From 3b34891e5b9c2694b7ffdc265290e25847dc3437 Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Thu, 19 Sep 2024 09:10:51 +0900 Subject: [PATCH 062/189] [SPARK-49684][CONNECT] Remove global locks from session and execution managers ### What changes were proposed in this pull request? Eliminate the use of global locks in the session and execution managers. Those locks residing in the streaming query manager cannot be easily removed because the tag and query maps seemingly need to be synchronised. ### Why are the changes needed? In order to achieve true scalability. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48131 from changgyoopark-db/SPARK-49684. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../SparkConnectExecutionManager.scala | 59 ++++++++---------- .../service/SparkConnectSessionManager.scala | 60 ++++++++----------- .../SparkConnectStreamingQueryCache.scala | 22 +++---- 3 files changed, 61 insertions(+), 80 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala index 61b41f932199e..d66964b8d34bd 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.connect.service import java.util.UUID import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit} -import javax.annotation.concurrent.GuardedBy +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} import scala.collection.mutable import scala.concurrent.duration.FiniteDuration @@ -66,7 +66,6 @@ private[connect] class SparkConnectExecutionManager() extends Logging { /** Concurrent hash table containing all the current executions. */ private val executions: ConcurrentMap[ExecuteKey, ExecuteHolder] = new ConcurrentHashMap[ExecuteKey, ExecuteHolder]() - private val executionsLock = new Object /** Graveyard of tombstones of executions that were abandoned and removed. */ private val abandonedTombstones = CacheBuilder @@ -74,13 +73,12 @@ private[connect] class SparkConnectExecutionManager() extends Logging { .maximumSize(SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_ABANDONED_TOMBSTONES_SIZE)) .build[ExecuteKey, ExecuteInfo]() - /** None if there are no executions. Otherwise, the time when the last execution was removed. */ - @GuardedBy("executionsLock") - private var lastExecutionTimeMs: Option[Long] = Some(System.currentTimeMillis()) + /** The time when the last execution was removed. */ + private var lastExecutionTimeMs: AtomicLong = new AtomicLong(System.currentTimeMillis()) /** Executor for the periodic maintenance */ - @GuardedBy("executionsLock") - private var scheduledExecutor: Option[ScheduledExecutorService] = None + private var scheduledExecutor: AtomicReference[ScheduledExecutorService] = + new AtomicReference[ScheduledExecutorService]() /** * Create a new ExecuteHolder and register it with this global manager and with its session. @@ -118,11 +116,6 @@ private[connect] class SparkConnectExecutionManager() extends Logging { sessionHolder.addExecuteHolder(executeHolder) - executionsLock.synchronized { - if (!executions.isEmpty()) { - lastExecutionTimeMs = None - } - } logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, executeHolder.key)} is created.") schedulePeriodicChecks() // Starts the maintenance thread if it hasn't started. @@ -151,11 +144,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { executions.remove(key) executeHolder.sessionHolder.removeExecuteHolder(executeHolder.operationId) - executionsLock.synchronized { - if (executions.isEmpty) { - lastExecutionTimeMs = Some(System.currentTimeMillis()) - } - } + updateLastExecutionTime() logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, key)} is removed.") @@ -197,7 +186,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { */ def listActiveExecutions: Either[Long, Seq[ExecuteInfo]] = { if (executions.isEmpty) { - Left(lastExecutionTimeMs.get) + Left(lastExecutionTimeMs.getAcquire()) } else { Right(executions.values().asScala.map(_.getExecuteInfo).toBuffer.toSeq) } @@ -212,22 +201,23 @@ private[connect] class SparkConnectExecutionManager() extends Logging { } private[connect] def shutdown(): Unit = { - executionsLock.synchronized { - scheduledExecutor.foreach { executor => - ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) - } - scheduledExecutor = None + val executor = scheduledExecutor.getAndSet(null) + if (executor != null) { + ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) } // note: this does not cleanly shut down the executions, but the server is shutting down. executions.clear() abandonedTombstones.invalidateAll() - executionsLock.synchronized { - if (lastExecutionTimeMs.isEmpty) { - lastExecutionTimeMs = Some(System.currentTimeMillis()) - } - } + updateLastExecutionTime() + } + + /** + * Updates the last execution time after the last execution has been removed. + */ + private def updateLastExecutionTime(): Unit = { + lastExecutionTimeMs.getAndUpdate(prev => prev.max(System.currentTimeMillis())) } /** @@ -235,16 +225,16 @@ private[connect] class SparkConnectExecutionManager() extends Logging { * for executions that have not been closed, but are left with no RPC attached to them, and * removes them after a timeout. */ - private def schedulePeriodicChecks(): Unit = executionsLock.synchronized { - scheduledExecutor match { - case Some(_) => // Already running. - case None => + private def schedulePeriodicChecks(): Unit = { + var executor = scheduledExecutor.getAcquire() + if (executor == null) { + executor = Executors.newSingleThreadScheduledExecutor() + if (scheduledExecutor.compareAndExchangeRelease(null, executor) == null) { val interval = SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_MAINTENANCE_INTERVAL) logInfo( log"Starting thread for cleanup of abandoned executions every " + log"${MDC(LogKeys.INTERVAL, interval)} ms") - scheduledExecutor = Some(Executors.newSingleThreadScheduledExecutor()) - scheduledExecutor.get.scheduleAtFixedRate( + executor.scheduleAtFixedRate( () => { try { val timeout = SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT) @@ -256,6 +246,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { interval, interval, TimeUnit.MILLISECONDS) + } } } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala index fec01813de6e2..4ca3a80bfb985 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.connect.service import java.util.UUID import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit} -import javax.annotation.concurrent.GuardedBy +import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable import scala.concurrent.duration.FiniteDuration @@ -40,8 +40,6 @@ import org.apache.spark.util.ThreadUtils */ class SparkConnectSessionManager extends Logging { - private val sessionsLock = new Object - private val sessionStore: ConcurrentMap[SessionKey, SessionHolder] = new ConcurrentHashMap[SessionKey, SessionHolder]() @@ -52,8 +50,8 @@ class SparkConnectSessionManager extends Logging { .build[SessionKey, SessionHolderInfo]() /** Executor for the periodic maintenance */ - @GuardedBy("sessionsLock") - private var scheduledExecutor: Option[ScheduledExecutorService] = None + private var scheduledExecutor: AtomicReference[ScheduledExecutorService] = + new AtomicReference[ScheduledExecutorService]() private def validateSessionId( key: SessionKey, @@ -75,8 +73,6 @@ class SparkConnectSessionManager extends Logging { val holder = getSession( key, Some(() => { - // Executed under sessionsState lock in getSession, to guard against concurrent removal - // and insertion into closedSessionsCache. validateSessionCreate(key) val holder = SessionHolder(key.userId, key.sessionId, newIsolatedSession()) holder.initializeSession() @@ -168,17 +164,14 @@ class SparkConnectSessionManager extends Logging { def closeSession(key: SessionKey): Unit = { val sessionHolder = removeSessionHolder(key) - // Rest of the cleanup outside sessionLock - the session cannot be accessed anymore by - // getOrCreateIsolatedSession. + // Rest of the cleanup: the session cannot be accessed anymore by getOrCreateIsolatedSession. sessionHolder.foreach(shutdownSessionHolder(_)) } private[connect] def shutdown(): Unit = { - sessionsLock.synchronized { - scheduledExecutor.foreach { executor => - ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) - } - scheduledExecutor = None + val executor = scheduledExecutor.getAndSet(null) + if (executor != null) { + ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) } // note: this does not cleanly shut down the sessions, but the server is shutting down. @@ -199,16 +192,16 @@ class SparkConnectSessionManager extends Logging { * * The checks are looking to remove sessions that expired. */ - private def schedulePeriodicChecks(): Unit = sessionsLock.synchronized { - scheduledExecutor match { - case Some(_) => // Already running. - case None => + private def schedulePeriodicChecks(): Unit = { + var executor = scheduledExecutor.getAcquire() + if (executor == null) { + executor = Executors.newSingleThreadScheduledExecutor() + if (scheduledExecutor.compareAndExchangeRelease(null, executor) == null) { val interval = SparkEnv.get.conf.get(CONNECT_SESSION_MANAGER_MAINTENANCE_INTERVAL) logInfo( log"Starting thread for cleanup of expired sessions every " + log"${MDC(INTERVAL, interval)} ms") - scheduledExecutor = Some(Executors.newSingleThreadScheduledExecutor()) - scheduledExecutor.get.scheduleAtFixedRate( + executor.scheduleAtFixedRate( () => { try { val defaultInactiveTimeoutMs = @@ -221,6 +214,7 @@ class SparkConnectSessionManager extends Logging { interval, interval, TimeUnit.MILLISECONDS) + } } } @@ -255,24 +249,18 @@ class SparkConnectSessionManager extends Logging { // .. and remove them. toRemove.foreach { sessionHolder => - // This doesn't use closeSession to be able to do the extra last chance check under lock. - val removedSession = { - // Last chance - check expiration time and remove under lock if expired. - val info = sessionHolder.getSessionHolderInfo - if (shouldExpire(info, System.currentTimeMillis())) { - logInfo( - log"Found session ${MDC(SESSION_HOLD_INFO, info)} that expired " + - log"and will be closed.") - removeSessionHolder(info.key) - } else { - None + val info = sessionHolder.getSessionHolderInfo + if (shouldExpire(info, System.currentTimeMillis())) { + logInfo( + log"Found session ${MDC(SESSION_HOLD_INFO, info)} that expired " + + log"and will be closed.") + removeSessionHolder(info.key) + try { + shutdownSessionHolder(sessionHolder) + } catch { + case NonFatal(ex) => logWarning("Unexpected exception closing session", ex) } } - // do shutdown and cleanup outside of lock. - try removedSession.foreach(shutdownSessionHolder(_)) - catch { - case NonFatal(ex) => logWarning("Unexpected exception closing session", ex) - } } logInfo("Finished periodic run of SparkConnectSessionManager maintenance.") } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala index 03719ddd87419..8241672d5107b 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql.connect.service import java.util.concurrent.Executors import java.util.concurrent.ScheduledExecutorService import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicReference import javax.annotation.concurrent.GuardedBy import scala.collection.mutable @@ -185,10 +186,10 @@ private[connect] class SparkConnectStreamingQueryCache( // Visible for testing. private[service] def shutdown(): Unit = queryCacheLock.synchronized { - scheduledExecutor.foreach { executor => + val executor = scheduledExecutor.getAndSet(null) + if (executor != null) { ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES)) } - scheduledExecutor = None } @GuardedBy("queryCacheLock") @@ -199,19 +200,19 @@ private[connect] class SparkConnectStreamingQueryCache( private val taggedQueries = new mutable.HashMap[String, mutable.ArrayBuffer[QueryCacheKey]] private val taggedQueriesLock = new Object - @GuardedBy("queryCacheLock") - private var scheduledExecutor: Option[ScheduledExecutorService] = None + private var scheduledExecutor: AtomicReference[ScheduledExecutorService] = + new AtomicReference[ScheduledExecutorService]() /** Schedules periodic checks if it is not already scheduled */ - private def schedulePeriodicChecks(): Unit = queryCacheLock.synchronized { - scheduledExecutor match { - case Some(_) => // Already running. - case None => + private def schedulePeriodicChecks(): Unit = { + var executor = scheduledExecutor.getAcquire() + if (executor == null) { + executor = Executors.newSingleThreadScheduledExecutor() + if (scheduledExecutor.compareAndExchangeRelease(null, executor) == null) { logInfo( log"Starting thread for polling streaming sessions " + log"every ${MDC(DURATION, sessionPollingPeriod.toMillis)}") - scheduledExecutor = Some(Executors.newSingleThreadScheduledExecutor()) - scheduledExecutor.get.scheduleAtFixedRate( + executor.scheduleAtFixedRate( () => { try periodicMaintenance() catch { @@ -221,6 +222,7 @@ private[connect] class SparkConnectStreamingQueryCache( sessionPollingPeriod.toMillis, sessionPollingPeriod.toMillis, TimeUnit.MILLISECONDS) + } } } From af45902d33c4d8e38a6427ac1d0c46fe057bb45a Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Wed, 18 Sep 2024 20:11:21 -0400 Subject: [PATCH 063/189] [SPARK-49422][CONNECT][SQL] Add groupByKey to sql/api ### What changes were proposed in this pull request? This PR adds `Dataset.groupByKey(..)` to the shared interface. I forgot to add in the previous PR. ### Why are the changes needed? The shared interface needs to support all functionality. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48147 from hvanhovell/SPARK-49422-follow-up. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../scala/org/apache/spark/sql/Dataset.scala | 24 ++----- .../org/apache/spark/sql/api/Dataset.scala | 22 ++++++ .../scala/org/apache/spark/sql/Dataset.scala | 68 +++---------------- 3 files changed, 39 insertions(+), 75 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index 161a0d9d265f0..accfff9f2b073 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -524,27 +524,11 @@ class Dataset[T] private[sql] ( result(0) } - /** - * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 3.5.0 - */ + /** @inheritdoc */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { KeyValueGroupedDatasetImpl[K, T](this, encoderFor[K], func) } - /** - * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 3.5.0 - */ - def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - groupByKey(ToScalaUDF(func))(encoder) - /** @inheritdoc */ @scala.annotation.varargs def rollup(cols: Column*): RelationalGroupedDataset = { @@ -1480,4 +1464,10 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) + + /** @inheritdoc */ + override def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala index 284a69fe6ee3e..7a3d6b0e03877 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala @@ -1422,6 +1422,28 @@ abstract class Dataset[T] extends Serializable { */ def reduce(func: ReduceFunction[T]): T = reduce(ToScalaUDF(func)) + /** + * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 2.0.0 + */ + def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T, DS] + + /** + * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 2.0.0 + */ + def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T, DS] = { + groupByKey(ToScalaUDF(func))(encoder) + } + /** * Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns * set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, except for the aggregation, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 61f9e6ff7c042..ef628ca612b49 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -62,7 +62,7 @@ import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, FileTable} import org.apache.spark.sql.execution.python.EvaluatePython import org.apache.spark.sql.execution.stat.StatFunctions -import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf, ToScalaUDF} +import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf} import org.apache.spark.sql.internal.ExpressionUtils.column import org.apache.spark.sql.internal.TypedAggUtils.withInputType import org.apache.spark.sql.streaming.DataStreamWriter @@ -865,24 +865,7 @@ class Dataset[T] private[sql]( Filter(condition.expr, logicalPlan) } - /** - * Groups the Dataset using the specified columns, so we can run aggregation on them. See - * [[RelationalGroupedDataset]] for all the available aggregate functions. - * - * {{{ - * // Compute the average for all numeric columns grouped by department. - * ds.groupBy($"department").avg() - * - * // Compute the max age and average salary, grouped by department and gender. - * ds.groupBy($"department", $"gender").agg(Map( - * "salary" -> "avg", - * "age" -> "max" - * )) - * }}} - * - * @group untypedrel - * @since 2.0.0 - */ + /** @inheritdoc */ @scala.annotation.varargs def groupBy(cols: Column*): RelationalGroupedDataset = { RelationalGroupedDataset(toDF(), cols.map(_.expr), RelationalGroupedDataset.GroupByType) @@ -914,13 +897,7 @@ class Dataset[T] private[sql]( rdd.reduce(func) } - /** - * (Scala-specific) - * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. - * - * @group typedrel - * @since 2.0.0 - */ + /** @inheritdoc */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { val withGroupingKey = AppendColumns(func, logicalPlan) val executed = sparkSession.sessionState.executePlan(withGroupingKey) @@ -933,16 +910,6 @@ class Dataset[T] private[sql]( withGroupingKey.newColumns) } - /** - * (Java-specific) - * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. - * - * @group typedrel - * @since 2.0.0 - */ - def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - groupByKey(ToScalaUDF(func))(encoder) - /** @inheritdoc */ def unpivot( ids: Array[Column], @@ -1640,28 +1607,7 @@ class Dataset[T] private[sql]( new DataFrameWriterV2Impl[T](table, this) } - /** - * Merges a set of updates, insertions, and deletions based on a source table into - * a target table. - * - * Scala Examples: - * {{{ - * spark.table("source") - * .mergeInto("target", $"source.id" === $"target.id") - * .whenMatched($"salary" === 100) - * .delete() - * .whenNotMatched() - * .insertAll() - * .whenNotMatchedBySource($"salary" === 100) - * .update(Map( - * "salary" -> lit(200) - * )) - * .merge() - * }}} - * - * @group basic - * @since 4.0.0 - */ + /** @inheritdoc */ def mergeInto(table: String, condition: Column): MergeIntoWriter[T] = { if (isStreaming) { logicalPlan.failAnalysis( @@ -2024,6 +1970,12 @@ class Dataset[T] private[sql]( @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) + /** @inheritdoc */ + override def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] + //////////////////////////////////////////////////////////////////////////// // For Python API //////////////////////////////////////////////////////////////////////////// From 58d73fe8e7cbff9878539d31430f819eff9fc7a1 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 19 Sep 2024 09:16:23 +0900 Subject: [PATCH 064/189] Revert "[SPARK-49495][DOCS][FOLLOWUP] Enable GitHub Pages settings via .asf.yml" This reverts commit b86e5d2ab1fb17f8dcbb5b4d50f3361494270438. --- .asf.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.asf.yaml b/.asf.yaml index 91a5f9b2bb1a2..22042b355b2fa 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -31,8 +31,6 @@ github: merge: false squash: true rebase: true - ghp_branch: master - ghp_path: /docs/_site notifications: pullrequests: reviews@spark.apache.org From 376382711e200aa978008b25630cc54271fd419b Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 19 Sep 2024 09:16:28 +0900 Subject: [PATCH 065/189] Revert "[SPARK-49495][DOCS][FOLLOWUP] Fix Pandoc installation for GitHub Pages publication action" This reverts commit 7de71a2ec78d985c2a045f13c1275101b126cec4. --- .github/workflows/pages.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index f10dadf315a1b..083620427c015 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -63,9 +63,9 @@ jobs: ruby-version: '3.3' bundler-cache: true - name: Install Pandoc - run: | - sudo apt-get update -y - sudo apt-get install pandoc + uses: pandoc/actions/setup@d6abb76f6c8a1a9a5e15a5190c96a02aabffd1ee + with: + version: 3.3 - name: Install dependencies for documentation generation run: | cd docs From 8861f0f9af3f397921ba1204cf4f76f4e20680bb Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 19 Sep 2024 09:16:33 +0900 Subject: [PATCH 066/189] Revert "[SPARK-49495][DOCS] Document and Feature Preview on the master branch via Live GitHub Pages Updates" This reverts commit b1807095bef9c6d98e60bdc2669c8af93bc68ad4. --- .github/workflows/pages.yml | 90 ------------------------------------- 1 file changed, 90 deletions(-) delete mode 100644 .github/workflows/pages.yml diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml deleted file mode 100644 index 083620427c015..0000000000000 --- a/.github/workflows/pages.yml +++ /dev/null @@ -1,90 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# - -name: GitHub Pages deployment - -on: - push: - branches: - - master - -concurrency: - group: 'docs preview' - cancel-in-progress: true - -jobs: - docs: - name: Build and deploy documentation - runs-on: ubuntu-latest - permissions: - id-token: write - pages: write - env: - SPARK_TESTING: 1 # Reduce some noise in the logs - RELEASE_VERSION: 'In-Progress' - steps: - - name: Checkout Spark repository - uses: actions/checkout@v4 - with: - repository: apache/spark - ref: 'master' - - name: Install Java 17 - uses: actions/setup-java@v4 - with: - distribution: zulu - java-version: 17 - - name: Install Python 3.9 - uses: actions/setup-python@v5 - with: - python-version: '3.9' - architecture: x64 - cache: 'pip' - - name: Install Python dependencies - run: pip install --upgrade -r dev/requirements.txt - - name: Install Ruby for documentation generation - uses: ruby/setup-ruby@v1 - with: - ruby-version: '3.3' - bundler-cache: true - - name: Install Pandoc - uses: pandoc/actions/setup@d6abb76f6c8a1a9a5e15a5190c96a02aabffd1ee - with: - version: 3.3 - - name: Install dependencies for documentation generation - run: | - cd docs - gem install bundler -v 2.4.22 -n /usr/local/bin - bundle install --retry=100 - - name: Run documentation build - run: | - sed -i".tmp1" 's/SPARK_VERSION:.*$/SPARK_VERSION: '"$RELEASE_VERSION"'/g' docs/_config.yml - sed -i".tmp2" 's/SPARK_VERSION_SHORT:.*$/SPARK_VERSION_SHORT: '"$RELEASE_VERSION"'/g' docs/_config.yml - sed -i".tmp3" "s/'facetFilters':.*$/'facetFilters': [\"version:$RELEASE_VERSION\"]/g" docs/_config.yml - sed -i".tmp4" 's/__version__: str = .*$/__version__: str = "'"$RELEASE_VERSION"'"/' python/pyspark/version.py - cd docs - SKIP_RDOC=1 bundle exec jekyll build - - name: Setup Pages - uses: actions/configure-pages@v5 - - name: Upload artifact - uses: actions/upload-pages-artifact@v3 - with: - path: 'docs/_site' - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 From f3c8d26eb0c3fd7f77950eb08c70bb2a9ab6493c Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 19 Sep 2024 10:36:03 +0900 Subject: [PATCH 067/189] Revert "[SPARK-49422][CONNECT][SQL] Add groupByKey to sql/api" This reverts commit af45902d33c4d8e38a6427ac1d0c46fe057bb45a. --- .../scala/org/apache/spark/sql/Dataset.scala | 24 +++++-- .../org/apache/spark/sql/api/Dataset.scala | 22 ------ .../scala/org/apache/spark/sql/Dataset.scala | 68 ++++++++++++++++--- 3 files changed, 75 insertions(+), 39 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index accfff9f2b073..161a0d9d265f0 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -524,11 +524,27 @@ class Dataset[T] private[sql] ( result(0) } - /** @inheritdoc */ + /** + * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 3.5.0 + */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { KeyValueGroupedDatasetImpl[K, T](this, encoderFor[K], func) } + /** + * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 3.5.0 + */ + def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + groupByKey(ToScalaUDF(func))(encoder) + /** @inheritdoc */ @scala.annotation.varargs def rollup(cols: Column*): RelationalGroupedDataset = { @@ -1464,10 +1480,4 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) - - /** @inheritdoc */ - override def groupByKey[K]( - func: MapFunction[T, K], - encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala index 7a3d6b0e03877..284a69fe6ee3e 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala @@ -1422,28 +1422,6 @@ abstract class Dataset[T] extends Serializable { */ def reduce(func: ReduceFunction[T]): T = reduce(ToScalaUDF(func)) - /** - * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 2.0.0 - */ - def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T, DS] - - /** - * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 2.0.0 - */ - def groupByKey[K]( - func: MapFunction[T, K], - encoder: Encoder[K]): KeyValueGroupedDataset[K, T, DS] = { - groupByKey(ToScalaUDF(func))(encoder) - } - /** * Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns * set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, except for the aggregation, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index ef628ca612b49..61f9e6ff7c042 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -62,7 +62,7 @@ import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, FileTable} import org.apache.spark.sql.execution.python.EvaluatePython import org.apache.spark.sql.execution.stat.StatFunctions -import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf} +import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf, ToScalaUDF} import org.apache.spark.sql.internal.ExpressionUtils.column import org.apache.spark.sql.internal.TypedAggUtils.withInputType import org.apache.spark.sql.streaming.DataStreamWriter @@ -865,7 +865,24 @@ class Dataset[T] private[sql]( Filter(condition.expr, logicalPlan) } - /** @inheritdoc */ + /** + * Groups the Dataset using the specified columns, so we can run aggregation on them. See + * [[RelationalGroupedDataset]] for all the available aggregate functions. + * + * {{{ + * // Compute the average for all numeric columns grouped by department. + * ds.groupBy($"department").avg() + * + * // Compute the max age and average salary, grouped by department and gender. + * ds.groupBy($"department", $"gender").agg(Map( + * "salary" -> "avg", + * "age" -> "max" + * )) + * }}} + * + * @group untypedrel + * @since 2.0.0 + */ @scala.annotation.varargs def groupBy(cols: Column*): RelationalGroupedDataset = { RelationalGroupedDataset(toDF(), cols.map(_.expr), RelationalGroupedDataset.GroupByType) @@ -897,7 +914,13 @@ class Dataset[T] private[sql]( rdd.reduce(func) } - /** @inheritdoc */ + /** + * (Scala-specific) + * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. + * + * @group typedrel + * @since 2.0.0 + */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { val withGroupingKey = AppendColumns(func, logicalPlan) val executed = sparkSession.sessionState.executePlan(withGroupingKey) @@ -910,6 +933,16 @@ class Dataset[T] private[sql]( withGroupingKey.newColumns) } + /** + * (Java-specific) + * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. + * + * @group typedrel + * @since 2.0.0 + */ + def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + groupByKey(ToScalaUDF(func))(encoder) + /** @inheritdoc */ def unpivot( ids: Array[Column], @@ -1607,7 +1640,28 @@ class Dataset[T] private[sql]( new DataFrameWriterV2Impl[T](table, this) } - /** @inheritdoc */ + /** + * Merges a set of updates, insertions, and deletions based on a source table into + * a target table. + * + * Scala Examples: + * {{{ + * spark.table("source") + * .mergeInto("target", $"source.id" === $"target.id") + * .whenMatched($"salary" === 100) + * .delete() + * .whenNotMatched() + * .insertAll() + * .whenNotMatchedBySource($"salary" === 100) + * .update(Map( + * "salary" -> lit(200) + * )) + * .merge() + * }}} + * + * @group basic + * @since 4.0.0 + */ def mergeInto(table: String, condition: Column): MergeIntoWriter[T] = { if (isStreaming) { logicalPlan.failAnalysis( @@ -1970,12 +2024,6 @@ class Dataset[T] private[sql]( @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) - /** @inheritdoc */ - override def groupByKey[K]( - func: MapFunction[T, K], - encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] - //////////////////////////////////////////////////////////////////////////// // For Python API //////////////////////////////////////////////////////////////////////////// From 3bdf146bbee58d207afaadc92024d9f6c4b941dd Mon Sep 17 00:00:00 2001 From: Mihailo Milosevic Date: Thu, 19 Sep 2024 09:27:38 +0200 Subject: [PATCH 068/189] [SPARK-49611][SQL][FOLLOW-UP] Fix wrong results of collations() TVF ### What changes were proposed in this pull request? Fix of accent sensitive and case sensitive column results. ### Why are the changes needed? When initial PR was introduced, ICU collation listing ended up with different order of generating columns so results were wrong. ### Does this PR introduce _any_ user-facing change? No, as spark 4.0 was not released yet. ### How was this patch tested? Existing test in CollationSuite.scala, which was wrong in the first place. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48152 from mihailom-db/tvf-collations-followup. Authored-by: Mihailo Milosevic Signed-off-by: Max Gekk --- .../sql/catalyst/util/CollationFactory.java | 4 ++-- .../org/apache/spark/sql/CollationSuite.scala | 24 +++++++++---------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java index 4b88e15e8ed72..87558971042e0 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java @@ -773,8 +773,8 @@ protected CollationMeta buildCollationMeta() { ICULocaleMap.get(locale).getDisplayCountry(), VersionInfo.ICU_VERSION.toString(), COLLATION_PAD_ATTRIBUTE, - caseSensitivity == CaseSensitivity.CS, - accentSensitivity == AccentSensitivity.AS); + accentSensitivity == AccentSensitivity.AS, + caseSensitivity == CaseSensitivity.CS); } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index d5d18b1ab081c..73fd897e91f53 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -1661,17 +1661,17 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { Row("SYSTEM", "BUILTIN", "UNICODE", "", "", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "UNICODE_AI", "", "", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "UNICODE_CI", "", "", "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "UNICODE_CI", "", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "UNICODE_CI_AI", "", "", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "af", "Afrikaans", "", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "af_AI", "Afrikaans", "", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "af_CI", "Afrikaans", "", "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "af_CI", "Afrikaans", "", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "af_CI_AI", "Afrikaans", "", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) @@ -1683,9 +1683,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { Seq(Row("SYSTEM", "BUILTIN", "zh_Hant_HKG", "Chinese", "Hong Kong SAR China", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_AI", "Chinese", "Hong Kong SAR China", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI", "Chinese", "Hong Kong SAR China", "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI", "Chinese", "Hong Kong SAR China", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hant_HKG_CI_AI", "Chinese", "Hong Kong SAR China", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) @@ -1693,9 +1693,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { Seq(Row("SYSTEM", "BUILTIN", "zh_Hans_SGP", "Chinese", "Singapore", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_AI", "Chinese", "Singapore", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_CI", "Chinese", "Singapore", "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_CI", "Chinese", "Singapore", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "zh_Hans_SGP_CI_AI", "Chinese", "Singapore", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) @@ -1704,17 +1704,17 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { Seq(Row("SYSTEM", "BUILTIN", "en_USA", "English", "United States", "ACCENT_SENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "en_USA_AI", "English", "United States", - "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), - Row("SYSTEM", "BUILTIN", "en_USA_CI", "English", "United States", "ACCENT_INSENSITIVE", "CASE_SENSITIVE", "NO_PAD", "75.1.0.0"), + Row("SYSTEM", "BUILTIN", "en_USA_CI", "English", "United States", + "ACCENT_SENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"), Row("SYSTEM", "BUILTIN", "en_USA_CI_AI", "English", "United States", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE", "NO_PAD", "75.1.0.0"))) checkAnswer(sql("SELECT NAME, LANGUAGE, ACCENT_SENSITIVITY, CASE_SENSITIVITY " + "FROM collations() WHERE COUNTRY = 'United States'"), Seq(Row("en_USA", "English", "ACCENT_SENSITIVE", "CASE_SENSITIVE"), - Row("en_USA_AI", "English", "ACCENT_SENSITIVE", "CASE_INSENSITIVE"), - Row("en_USA_CI", "English", "ACCENT_INSENSITIVE", "CASE_SENSITIVE"), + Row("en_USA_AI", "English", "ACCENT_INSENSITIVE", "CASE_SENSITIVE"), + Row("en_USA_CI", "English", "ACCENT_SENSITIVE", "CASE_INSENSITIVE"), Row("en_USA_CI_AI", "English", "ACCENT_INSENSITIVE", "CASE_INSENSITIVE"))) checkAnswer(sql("SELECT NAME FROM collations() WHERE ICU_VERSION is null"), From 492d1b14c0d19fa89b9ce9c0e48fc0e4c120b70c Mon Sep 17 00:00:00 2001 From: Anton Okolnychyi Date: Thu, 19 Sep 2024 11:09:40 +0200 Subject: [PATCH 069/189] [SPARK-48782][SQL] Add support for executing procedures in catalogs ### What changes were proposed in this pull request? This PR adds support for executing procedures in catalogs. ### Why are the changes needed? These changes are needed per [discussed and voted](https://lists.apache.org/thread/w586jr53fxwk4pt9m94b413xyjr1v25m) SPIP tracked in [SPARK-44167](https://issues.apache.org/jira/browse/SPARK-44167). ### Does this PR introduce _any_ user-facing change? Yes. This PR adds CALL commands. ### How was this patch tested? This PR comes with tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47943 from aokolnychyi/spark-48782. Authored-by: Anton Okolnychyi Signed-off-by: Wenchen Fan --- .../resources/error/error-conditions.json | 6 + docs/sql-ref-ansi-compliance.md | 1 + .../spark/sql/catalyst/parser/SqlBaseLexer.g4 | 1 + .../sql/catalyst/parser/SqlBaseParser.g4 | 5 + .../procedures/ProcedureParameter.java | 5 + .../catalog/procedures/UnboundProcedure.java | 6 + .../sql/catalyst/analysis/Analyzer.scala | 65 +- .../catalyst/analysis/AnsiTypeCoercion.scala | 1 + .../sql/catalyst/analysis/CheckAnalysis.scala | 8 + .../sql/catalyst/analysis/TypeCoercion.scala | 16 + .../spark/sql/catalyst/analysis/package.scala | 6 +- .../catalyst/analysis/v2ResolutionPlans.scala | 17 +- .../sql/catalyst/parser/AstBuilder.scala | 22 + .../logical/ExecutableDuringAnalysis.scala | 28 + .../plans/logical/FunctionBuilderBase.scala | 36 +- .../catalyst/plans/logical/MultiResult.scala | 30 + .../catalyst/plans/logical/v2Commands.scala | 67 +- .../sql/catalyst/rules/RuleIdCollection.scala | 1 + .../sql/catalyst/trees/TreePatterns.scala | 1 + .../catalog/CatalogV2Implicits.scala | 7 + .../sql/errors/QueryCompilationErrors.scala | 7 + .../connector/catalog/InMemoryCatalog.scala | 19 +- .../catalyst/analysis/InvokeProcedures.scala | 71 ++ .../spark/sql/execution/MultiResultExec.scala | 36 + .../spark/sql/execution/SparkStrategies.scala | 2 + .../sql/execution/command/commands.scala | 11 +- .../datasources/v2/DataSourceV2Strategy.scala | 6 +- .../datasources/v2/ExplainOnlySparkPlan.scala | 38 + .../internal/BaseSessionStateBuilder.scala | 3 +- .../sql-tests/results/ansi/keywords.sql.out | 2 + .../sql-tests/results/keywords.sql.out | 1 + .../spark/sql/connector/ProcedureSuite.scala | 654 ++++++++++++++++++ .../ThriftServerWithSparkContextSuite.scala | 2 +- .../sql/hive/HiveSessionStateBuilder.scala | 3 +- 34 files changed, 1162 insertions(+), 22 deletions(-) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ExecutableDuringAnalysis.scala create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/MultiResult.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/InvokeProcedures.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/MultiResultExec.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExplainOnlySparkPlan.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 6463cc2c12da7..72985de6631f0 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1456,6 +1456,12 @@ ], "sqlState" : "2203G" }, + "FAILED_TO_LOAD_ROUTINE" : { + "message" : [ + "Failed to load routine ." + ], + "sqlState" : "38000" + }, "FAILED_TO_PARSE_TOO_COMPLEX" : { "message" : [ "The statement, including potential SQL functions and referenced views, was too complex to parse.", diff --git a/docs/sql-ref-ansi-compliance.md b/docs/sql-ref-ansi-compliance.md index fff6906457f7d..12dff1e325c49 100644 --- a/docs/sql-ref-ansi-compliance.md +++ b/docs/sql-ref-ansi-compliance.md @@ -426,6 +426,7 @@ Below is a list of all the keywords in Spark SQL. |BY|non-reserved|non-reserved|reserved| |BYTE|non-reserved|non-reserved|non-reserved| |CACHE|non-reserved|non-reserved|non-reserved| +|CALL|reserved|non-reserved|reserved| |CALLED|non-reserved|non-reserved|non-reserved| |CASCADE|non-reserved|non-reserved|non-reserved| |CASE|reserved|non-reserved|reserved| diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index e704f9f58b964..de28041acd41f 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -146,6 +146,7 @@ BUCKETS: 'BUCKETS'; BY: 'BY'; BYTE: 'BYTE'; CACHE: 'CACHE'; +CALL: 'CALL'; CALLED: 'CALLED'; CASCADE: 'CASCADE'; CASE: 'CASE'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index f13dde773496a..e591a43b84d1a 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -298,6 +298,10 @@ statement LEFT_PAREN columns=multipartIdentifierPropertyList RIGHT_PAREN (OPTIONS options=propertyList)? #createIndex | DROP INDEX (IF EXISTS)? identifier ON TABLE? identifierReference #dropIndex + | CALL identifierReference + LEFT_PAREN + (functionArgument (COMMA functionArgument)*)? + RIGHT_PAREN #call | unsupportedHiveNativeCommands .*? #failNativeCommand ; @@ -1851,6 +1855,7 @@ nonReserved | BY | BYTE | CACHE + | CALL | CALLED | CASCADE | CASE diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/ProcedureParameter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/ProcedureParameter.java index 90d531ae21892..18c76833c5879 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/ProcedureParameter.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/ProcedureParameter.java @@ -32,6 +32,11 @@ */ @Evolving public interface ProcedureParameter { + /** + * A field metadata key that indicates whether an argument is passed by name. + */ + String BY_NAME_METADATA_KEY = "BY_NAME"; + /** * Creates a builder for an IN procedure parameter. * diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/UnboundProcedure.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/UnboundProcedure.java index ee9a09055243b..1a91fd21bf07e 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/UnboundProcedure.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/procedures/UnboundProcedure.java @@ -35,6 +35,12 @@ public interface UnboundProcedure extends Procedure { * validate if the input types are compatible while binding or delegate that to Spark. Regardless, * Spark will always perform the final validation of the arguments and rearrange them as needed * based on {@link BoundProcedure#parameters() reported parameters}. + *

    + * The provided {@code inputType} is based on the procedure arguments. If an argument is passed + * by name, its metadata will indicate this with {@link ProcedureParameter#BY_NAME_METADATA_KEY} + * set to {@code true}. In such cases, the field name will match the name of the target procedure + * parameter. If the argument is not named, {@link ProcedureParameter#BY_NAME_METADATA_KEY} will + * not be set and the name will be assigned randomly. * * @param inputType the input types to bind to * @return the bound procedure that is most suitable for the given input types diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index 0164af945ca28..9e5b1d1254c87 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala @@ -26,7 +26,7 @@ import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.util.{Failure, Random, Success, Try} -import org.apache.spark.{SparkException, SparkUnsupportedOperationException} +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst._ import org.apache.spark.sql.catalyst.catalog._ @@ -50,6 +50,7 @@ import org.apache.spark.sql.connector.catalog.{View => _, _} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ import org.apache.spark.sql.connector.catalog.TableChange.{After, ColumnPosition} import org.apache.spark.sql.connector.catalog.functions.{AggregateFunction => V2AggregateFunction, ScalarFunction, UnboundFunction} +import org.apache.spark.sql.connector.catalog.procedures.{BoundProcedure, ProcedureParameter, UnboundProcedure} import org.apache.spark.sql.connector.expressions.{FieldReference, IdentityTransform} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation @@ -310,6 +311,8 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor ExtractGenerator :: ResolveGenerate :: ResolveFunctions :: + ResolveProcedures :: + BindProcedures :: ResolveTableSpec :: ResolveAliases :: ResolveSubquery :: @@ -2611,6 +2614,66 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor } } + /** + * A rule that resolves procedures. + */ + object ResolveProcedures extends Rule[LogicalPlan] { + def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(UNRESOLVED_PROCEDURE), ruleId) { + case Call(UnresolvedProcedure(CatalogAndIdentifier(catalog, ident)), args, execute) => + val procedureCatalog = catalog.asProcedureCatalog + val procedure = load(procedureCatalog, ident) + Call(ResolvedProcedure(procedureCatalog, ident, procedure), args, execute) + } + + private def load(catalog: ProcedureCatalog, ident: Identifier): UnboundProcedure = { + try { + catalog.loadProcedure(ident) + } catch { + case e: Exception if !e.isInstanceOf[SparkThrowable] => + val nameParts = catalog.name +: ident.asMultipartIdentifier + throw QueryCompilationErrors.failedToLoadRoutineError(nameParts, e) + } + } + } + + /** + * A rule that binds procedures to the input types and rearranges arguments as needed. + */ + object BindProcedures extends Rule[LogicalPlan] { + def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + case Call(ResolvedProcedure(catalog, ident, unbound: UnboundProcedure), args, execute) + if args.forall(_.resolved) => + val inputType = extractInputType(args) + val bound = unbound.bind(inputType) + validateParameterModes(bound) + val rearrangedArgs = NamedParametersSupport.defaultRearrange(bound, args) + Call(ResolvedProcedure(catalog, ident, bound), rearrangedArgs, execute) + } + + private def extractInputType(args: Seq[Expression]): StructType = { + val fields = args.zipWithIndex.map { + case (NamedArgumentExpression(name, value), _) => + StructField(name, value.dataType, value.nullable, byNameMetadata) + case (arg, index) => + StructField(s"param$index", arg.dataType, arg.nullable) + } + StructType(fields) + } + + private def byNameMetadata: Metadata = { + new MetadataBuilder() + .putBoolean(ProcedureParameter.BY_NAME_METADATA_KEY, value = true) + .build() + } + + private def validateParameterModes(procedure: BoundProcedure): Unit = { + procedure.parameters.find(_.mode != ProcedureParameter.Mode.IN).foreach { param => + throw SparkException.internalError(s"Unsupported parameter mode: ${param.mode}") + } + } + } + /** * This rule resolves and rewrites subqueries inside expressions. * diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala index 17b1c4e249f57..3afe0ec8e9a7a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala @@ -77,6 +77,7 @@ object AnsiTypeCoercion extends TypeCoercionBase { override def typeCoercionRules: List[Rule[LogicalPlan]] = UnpivotCoercion :: WidenSetOperationTypes :: + ProcedureArgumentCoercion :: new AnsiCombinedTypeCoercionRule( CollationTypeCasts :: InConversion :: diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala index 752ff49e1f90d..5a9d5cd87ecc7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala @@ -676,6 +676,14 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB varName, c.defaultExpr.originalSQL) + case c: Call if c.resolved && c.bound && c.checkArgTypes().isFailure => + c.checkArgTypes() match { + case mismatch: TypeCheckResult.DataTypeMismatch => + c.dataTypeMismatch("CALL", mismatch) + case _ => + throw SparkException.internalError("Invalid input for procedure") + } + case _ => // Falls back to the following checks } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala index 08c5b3531b4c8..5983346ff1e27 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.AlwaysProcess import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.procedures.BoundProcedure import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.types.{AbstractArrayType, AbstractMapType, AbstractStringType, StringTypeAnyCollation} @@ -202,6 +203,20 @@ abstract class TypeCoercionBase { } } + /** + * A type coercion rule that implicitly casts procedure arguments to expected types. + */ + object ProcedureArgumentCoercion extends Rule[LogicalPlan] { + override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + case c @ Call(ResolvedProcedure(_, _, procedure: BoundProcedure), args, _) if c.resolved => + val expectedDataTypes = procedure.parameters.map(_.dataType) + val coercedArgs = args.zip(expectedDataTypes).map { + case (arg, expectedType) => implicitCast(arg, expectedType).getOrElse(arg) + } + c.copy(args = coercedArgs) + } + } + /** * Widens the data types of the [[Unpivot]] values. */ @@ -838,6 +853,7 @@ object TypeCoercion extends TypeCoercionBase { override def typeCoercionRules: List[Rule[LogicalPlan]] = UnpivotCoercion :: WidenSetOperationTypes :: + ProcedureArgumentCoercion :: new CombinedTypeCoercionRule( CollationTypeCasts :: InConversion :: diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/package.scala index c0689eb121679..daab9e4d78bf5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/package.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/package.scala @@ -67,9 +67,13 @@ package object analysis { } def dataTypeMismatch(expr: Expression, mismatch: DataTypeMismatch): Nothing = { + dataTypeMismatch(toSQLExpr(expr), mismatch) + } + + def dataTypeMismatch(sqlExpr: String, mismatch: DataTypeMismatch): Nothing = { throw new AnalysisException( errorClass = s"DATATYPE_MISMATCH.${mismatch.errorSubClass}", - messageParameters = mismatch.messageParameters + ("sqlExpr" -> toSQLExpr(expr)), + messageParameters = mismatch.messageParameters + ("sqlExpr" -> sqlExpr), origin = t.origin) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/v2ResolutionPlans.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/v2ResolutionPlans.scala index ecdf40e87a894..dee78b8f03af4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/v2ResolutionPlans.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/v2ResolutionPlans.scala @@ -23,13 +23,14 @@ import org.apache.spark.sql.catalyst.catalog.CatalogTable import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec import org.apache.spark.sql.catalyst.expressions.{Attribute, LeafExpression, Unevaluable} import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, Statistics} -import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern, UNRESOLVED_FUNC} +import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern, UNRESOLVED_FUNC, UNRESOLVED_PROCEDURE} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CharVarcharUtils -import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, Table, TableCatalog} +import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, ProcedureCatalog, Table, TableCatalog} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ import org.apache.spark.sql.connector.catalog.TableChange.ColumnPosition import org.apache.spark.sql.connector.catalog.functions.UnboundFunction +import org.apache.spark.sql.connector.catalog.procedures.Procedure import org.apache.spark.sql.types.{DataType, StructField} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -135,6 +136,12 @@ case class UnresolvedFunctionName( case class UnresolvedIdentifier(nameParts: Seq[String], allowTemp: Boolean = false) extends UnresolvedLeafNode +/** + * A procedure identifier that should be resolved into [[ResolvedProcedure]]. + */ +case class UnresolvedProcedure(nameParts: Seq[String]) extends UnresolvedLeafNode { + final override val nodePatterns: Seq[TreePattern] = Seq(UNRESOLVED_PROCEDURE) +} /** * A resolved leaf node whose statistics has no meaning. @@ -192,6 +199,12 @@ case class ResolvedFieldName(path: Seq[String], field: StructField) extends Fiel case class ResolvedFieldPosition(position: ColumnPosition) extends FieldPosition +case class ResolvedProcedure( + catalog: ProcedureCatalog, + ident: Identifier, + procedure: Procedure) extends LeafNodeWithoutStats { + override def output: Seq[Attribute] = Nil +} /** * A plan containing resolved persistent views. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index cb0e0e35c3704..52529bb4b789b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -5697,6 +5697,28 @@ class AstBuilder extends DataTypeAstBuilder ctx.EXISTS != null) } + /** + * Creates a plan for invoking a procedure. + * + * For example: + * {{{ + * CALL multi_part_name(v1, v2, ...); + * CALL multi_part_name(v1, param2 => v2, ...); + * CALL multi_part_name(param1 => v1, param2 => v2, ...); + * }}} + */ + override def visitCall(ctx: CallContext): LogicalPlan = withOrigin(ctx) { + val procedure = withIdentClause(ctx.identifierReference, UnresolvedProcedure) + val args = ctx.functionArgument.asScala.map { + case expr if expr.namedArgumentExpression != null => + val namedExpr = expr.namedArgumentExpression + NamedArgumentExpression(namedExpr.key.getText, expression(namedExpr.value)) + case expr => + expression(expr) + }.toSeq + Call(procedure, args) + } + /** * Create a TimestampAdd expression. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ExecutableDuringAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ExecutableDuringAnalysis.scala new file mode 100644 index 0000000000000..dc8dbf701f6a9 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ExecutableDuringAnalysis.scala @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.plans.logical + +/** + * A logical plan node that requires execution during analysis. + */ +trait ExecutableDuringAnalysis extends LogicalPlan { + /** + * Returns the logical plan node that should be used for EXPLAIN. + */ + def stageForExplain(): LogicalPlan +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/FunctionBuilderBase.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/FunctionBuilderBase.scala index 4701f4ea1e172..75b2fcd3a5f34 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/FunctionBuilderBase.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/FunctionBuilderBase.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.catalyst.plans.logical import org.apache.spark.sql.catalyst.expressions.{Expression, NamedArgumentExpression} +import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns +import org.apache.spark.sql.connector.catalog.procedures.{BoundProcedure, ProcedureParameter} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.util.ArrayImplicits._ @@ -122,12 +124,32 @@ object NamedParametersSupport { functionSignature: FunctionSignature, args: Seq[Expression], functionName: String): Seq[Expression] = { - val parameters: Seq[InputParameter] = functionSignature.parameters + defaultRearrange(functionName, functionSignature.parameters, args) + } + + final def defaultRearrange(procedure: BoundProcedure, args: Seq[Expression]): Seq[Expression] = { + defaultRearrange( + procedure.name, + procedure.parameters.map(toInputParameter).toSeq, + args) + } + + private def toInputParameter(param: ProcedureParameter): InputParameter = { + val defaultValue = Option(param.defaultValueExpression).map { expr => + ResolveDefaultColumns.analyze(param.name, param.dataType, expr, "CALL") + } + InputParameter(param.name, defaultValue) + } + + private def defaultRearrange( + routineName: String, + parameters: Seq[InputParameter], + args: Seq[Expression]): Seq[Expression] = { if (parameters.dropWhile(_.default.isEmpty).exists(_.default.isEmpty)) { - throw QueryCompilationErrors.unexpectedRequiredParameter(functionName, parameters) + throw QueryCompilationErrors.unexpectedRequiredParameter(routineName, parameters) } - val (positionalArgs, namedArgs) = splitAndCheckNamedArguments(args, functionName) + val (positionalArgs, namedArgs) = splitAndCheckNamedArguments(args, routineName) val namedParameters: Seq[InputParameter] = parameters.drop(positionalArgs.size) // The following loop checks for the following: @@ -140,12 +162,12 @@ object NamedParametersSupport { namedArgs.foreach { namedArg => val parameterName = namedArg.key if (!parameterNamesSet.contains(parameterName)) { - throw QueryCompilationErrors.unrecognizedParameterName(functionName, namedArg.key, + throw QueryCompilationErrors.unrecognizedParameterName(routineName, namedArg.key, parameterNamesSet.toSeq) } if (positionalParametersSet.contains(parameterName)) { throw QueryCompilationErrors.positionalAndNamedArgumentDoubleReference( - functionName, namedArg.key) + routineName, namedArg.key) } } @@ -154,7 +176,7 @@ object NamedParametersSupport { val validParameterSizes = Array.range(parameters.count(_.default.isEmpty), parameters.size + 1).toImmutableArraySeq throw QueryCompilationErrors.wrongNumArgsError( - functionName, validParameterSizes, args.length) + routineName, validParameterSizes, args.length) } // This constructs a map from argument name to value for argument rearrangement. @@ -168,7 +190,7 @@ object NamedParametersSupport { namedArgMap.getOrElse( param.name, if (param.default.isEmpty) { - throw QueryCompilationErrors.requiredParameterNotFound(functionName, param.name, index) + throw QueryCompilationErrors.requiredParameterNotFound(routineName, param.name, index) } else { param.default.get } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/MultiResult.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/MultiResult.scala new file mode 100644 index 0000000000000..f249e5c87eba2 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/MultiResult.scala @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.plans.logical + +import org.apache.spark.sql.catalyst.expressions.Attribute + +case class MultiResult(children: Seq[LogicalPlan]) extends LogicalPlan { + + override def output: Seq[Attribute] = children.lastOption.map(_.output).getOrElse(Nil) + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[LogicalPlan]): MultiResult = { + copy(children = newChildren) + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala index fdd43404e1d98..b465e0e11612f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala @@ -19,17 +19,22 @@ package org.apache.spark.sql.catalyst.plans.logical import org.apache.spark.{SparkIllegalArgumentException, SparkUnsupportedOperationException} import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AssignmentUtils, EliminateSubqueryAliases, FieldName, NamedRelation, PartitionSpec, ResolvedIdentifier, UnresolvedException, ViewSchemaMode} +import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AssignmentUtils, EliminateSubqueryAliases, FieldName, NamedRelation, PartitionSpec, ResolvedIdentifier, ResolvedProcedure, TypeCheckResult, UnresolvedException, UnresolvedProcedure, ViewSchemaMode} +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess} import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec import org.apache.spark.sql.catalyst.catalog.FunctionResource import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, MetadataAttribute, NamedExpression, UnaryExpression, Unevaluable, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.DescribeCommandSchema import org.apache.spark.sql.catalyst.trees.BinaryLike -import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, RowDeltaUtils, WriteDeltaProjections} +import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, truncatedString, CharVarcharUtils, RowDeltaUtils, WriteDeltaProjections} +import org.apache.spark.sql.catalyst.util.TypeUtils.{ordinalNumber, toSQLExpr} import org.apache.spark.sql.connector.catalog._ +import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.{IdentifierHelper, MultipartIdentifierHelper} +import org.apache.spark.sql.connector.catalog.procedures.BoundProcedure import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.write.{DeltaWrite, RowLevelOperation, RowLevelOperationTable, SupportsDelta, Write} +import org.apache.spark.sql.errors.DataTypeErrors.toSQLType import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, MapType, MetadataBuilder, StringType, StructField, StructType} import org.apache.spark.util.ArrayImplicits._ @@ -1571,3 +1576,61 @@ case class SetVariable( override protected def withNewChildInternal(newChild: LogicalPlan): SetVariable = copy(sourceQuery = newChild) } + +/** + * The logical plan of the CALL statement. + */ +case class Call( + procedure: LogicalPlan, + args: Seq[Expression], + execute: Boolean = true) + extends UnaryNode with ExecutableDuringAnalysis { + + override def output: Seq[Attribute] = Nil + + override def child: LogicalPlan = procedure + + def bound: Boolean = procedure match { + case ResolvedProcedure(_, _, _: BoundProcedure) => true + case _ => false + } + + def checkArgTypes(): TypeCheckResult = { + require(resolved && bound, "can check arg types only after resolution and binding") + + val params = procedure match { + case ResolvedProcedure(_, _, bound: BoundProcedure) => bound.parameters + } + require(args.length == params.length, "number of args and params must match after binding") + + args.zip(params).zipWithIndex.collectFirst { + case ((arg, param), idx) + if !DataType.equalsIgnoreCompatibleNullability(arg.dataType, param.dataType) => + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(idx), + "requiredType" -> toSQLType(param.dataType), + "inputSql" -> toSQLExpr(arg), + "inputType" -> toSQLType(arg.dataType))) + }.getOrElse(TypeCheckSuccess) + } + + override def simpleString(maxFields: Int): String = { + val name = procedure match { + case ResolvedProcedure(catalog, ident, _) => + s"${quoteIfNeeded(catalog.name)}.${ident.quoted}" + case UnresolvedProcedure(nameParts) => + nameParts.quoted + } + val argsString = truncatedString(args, ", ", maxFields) + s"Call $name($argsString)" + } + + override def stageForExplain(): Call = { + copy(execute = false) + } + + override protected def withNewChildInternal(newChild: LogicalPlan): Call = + copy(procedure = newChild) +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala index c70b43f0db173..b5556cbae7cd9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala @@ -54,6 +54,7 @@ object RuleIdCollection { "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveDeserializer" :: "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveEncodersInUDF" :: "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveFunctions" :: + "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveProcedures" :: "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveGenerate" :: "org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveGroupingAnalytics" :: "org.apache.spark.sql.catalyst.analysis.ResolveHigherOrderFunctions" :: diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala index 826ac52c2b817..0f1c98b53e0b3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala @@ -157,6 +157,7 @@ object TreePattern extends Enumeration { // Unresolved Plan patterns (Alphabetically ordered) val UNRESOLVED_FUNC: Value = Value + val UNRESOLVED_PROCEDURE: Value = Value val UNRESOLVED_SUBQUERY_COLUMN_ALIAS: Value = Value val UNRESOLVED_TABLE_VALUED_FUNCTION: Value = Value val UNRESOLVED_TRANSPOSE: Value = Value diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala index 65bdae85be12a..282350dda67d3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala @@ -126,6 +126,13 @@ private[sql] object CatalogV2Implicits { case _ => throw QueryCompilationErrors.missingCatalogAbilityError(plugin, "functions") } + + def asProcedureCatalog: ProcedureCatalog = plugin match { + case procedureCatalog: ProcedureCatalog => + procedureCatalog + case _ => + throw QueryCompilationErrors.missingCatalogAbilityError(plugin, "procedures") + } } implicit class NamespaceHelper(namespace: Array[String]) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index ad0e1d07bf93d..0b5255e95f073 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -853,6 +853,13 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat origin = origin) } + def failedToLoadRoutineError(nameParts: Seq[String], e: Exception): Throwable = { + new AnalysisException( + errorClass = "FAILED_TO_LOAD_ROUTINE", + messageParameters = Map("routineName" -> toSQLId(nameParts)), + cause = Some(e)) + } + def unresolvedRoutineError(name: FunctionIdentifier, searchPath: Seq[String]): Throwable = { new AnalysisException( errorClass = "UNRESOLVED_ROUTINE", diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalog.scala index 8d8d2317f0986..411a88b8765f6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalog.scala @@ -24,10 +24,13 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.analysis.{NoSuchFunctionException, NoSuchNamespaceException} import org.apache.spark.sql.connector.catalog.functions.UnboundFunction +import org.apache.spark.sql.connector.catalog.procedures.UnboundProcedure -class InMemoryCatalog extends InMemoryTableCatalog with FunctionCatalog { +class InMemoryCatalog extends InMemoryTableCatalog with FunctionCatalog with ProcedureCatalog { protected val functions: util.Map[Identifier, UnboundFunction] = new ConcurrentHashMap[Identifier, UnboundFunction]() + protected val procedures: util.Map[Identifier, UnboundProcedure] = + new ConcurrentHashMap[Identifier, UnboundProcedure]() override protected def allNamespaces: Seq[Seq[String]] = { (tables.keySet.asScala.map(_.namespace.toSeq) ++ @@ -63,4 +66,18 @@ class InMemoryCatalog extends InMemoryTableCatalog with FunctionCatalog { def clearFunctions(): Unit = { functions.clear() } + + override def loadProcedure(ident: Identifier): UnboundProcedure = { + val procedure = procedures.get(ident) + if (procedure == null) throw new RuntimeException("Procedure not found: " + ident) + procedure + } + + def createProcedure(ident: Identifier, procedure: UnboundProcedure): UnboundProcedure = { + procedures.put(ident, procedure) + } + + def clearProcedures(): Unit = { + procedures.clear() + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/InvokeProcedures.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/InvokeProcedures.scala new file mode 100644 index 0000000000000..c7320d350a7ff --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/InvokeProcedures.scala @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.analysis + +import scala.jdk.CollectionConverters.IteratorHasAsScala + +import org.apache.spark.SparkException +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow} +import org.apache.spark.sql.catalyst.plans.logical.{Call, LocalRelation, LogicalPlan, MultiResult} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.procedures.BoundProcedure +import org.apache.spark.sql.connector.read.{LocalScan, Scan} +import org.apache.spark.util.ArrayImplicits._ + +class InvokeProcedures(session: SparkSession) extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + case c: Call if c.resolved && c.bound && c.execute && c.checkArgTypes().isSuccess => + session.sessionState.optimizer.execute(c) match { + case Call(ResolvedProcedure(_, _, procedure: BoundProcedure), args, _) => + invoke(procedure, args) + case _ => + throw SparkException.internalError("Unexpected plan for optimized CALL statement") + } + } + + private def invoke(procedure: BoundProcedure, args: Seq[Expression]): LogicalPlan = { + val input = toInternalRow(args) + val scanIterator = procedure.call(input) + val relations = scanIterator.asScala.map(toRelation).toSeq + relations match { + case Nil => LocalRelation(Nil) + case Seq(relation) => relation + case _ => MultiResult(relations) + } + } + + private def toRelation(scan: Scan): LogicalPlan = scan match { + case s: LocalScan => + val attrs = DataTypeUtils.toAttributes(s.readSchema) + val data = s.rows.toImmutableArraySeq + LocalRelation(attrs, data) + case _ => + throw SparkException.internalError( + s"Only local scans are temporarily supported as procedure output: ${scan.getClass.getName}") + } + + private def toInternalRow(args: Seq[Expression]): InternalRow = { + require(args.forall(_.foldable), "args must be foldable") + val values = args.map(_.eval()).toArray + new GenericInternalRow(values) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/MultiResultExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/MultiResultExec.scala new file mode 100644 index 0000000000000..c2b12b053c927 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/MultiResultExec.scala @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute + +case class MultiResultExec(children: Seq[SparkPlan]) extends SparkPlan { + + override def output: Seq[Attribute] = children.lastOption.map(_.output).getOrElse(Nil) + + override protected def doExecute(): RDD[InternalRow] = { + children.lastOption.map(_.execute()).getOrElse(sparkContext.emptyRDD) + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[SparkPlan]): MultiResultExec = { + copy(children = newChildren) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala index 6d940a30619fb..aee735e48fc5c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala @@ -1041,6 +1041,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { case WriteFiles(child, fileFormat, partitionColumns, bucket, options, staticPartitions) => WriteFilesExec(planLater(child), fileFormat, partitionColumns, bucket, options, staticPartitions) :: Nil + case MultiResult(children) => + MultiResultExec(children.map(planLater)) :: Nil case _ => Nil } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/commands.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/commands.scala index ea2736b2c1266..ea9d53190546e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/commands.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/commands.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.{Row, SparkSession} import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.QueryPlan -import org.apache.spark.sql.catalyst.plans.logical.{Command, LogicalPlan, SupervisingCommand} +import org.apache.spark.sql.catalyst.plans.logical.{Command, ExecutableDuringAnalysis, LogicalPlan, SupervisingCommand} import org.apache.spark.sql.catalyst.trees.{LeafLike, UnaryLike} import org.apache.spark.sql.connector.ExternalCommandRunner import org.apache.spark.sql.execution.{CommandExecutionMode, ExplainMode, LeafExecNode, SparkPlan, UnaryExecNode} @@ -165,14 +165,19 @@ case class ExplainCommand( // Run through the optimizer to generate the physical plan. override def run(sparkSession: SparkSession): Seq[Row] = try { - val outputString = sparkSession.sessionState.executePlan(logicalPlan, CommandExecutionMode.SKIP) - .explainString(mode) + val stagedLogicalPlan = stageForAnalysis(logicalPlan) + val qe = sparkSession.sessionState.executePlan(stagedLogicalPlan, CommandExecutionMode.SKIP) + val outputString = qe.explainString(mode) Seq(Row(outputString)) } catch { case NonFatal(cause) => ("Error occurred during query planning: \n" + cause.getMessage).split("\n") .map(Row(_)).toImmutableArraySeq } + private def stageForAnalysis(plan: LogicalPlan): LogicalPlan = plan transform { + case p: ExecutableDuringAnalysis => p.stageForExplain() + } + def withTransformedSupervisedPlan(transformer: LogicalPlan => LogicalPlan): LogicalPlan = copy(logicalPlan = transformer(logicalPlan)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index d7f46c32f99a0..76cd33b815edd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -32,8 +32,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, DynamicPruning import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.util.{toPrettySQL, GeneratedColumn, - IdentityColumn, ResolveDefaultColumns, V2ExpressionBuilder} +import org.apache.spark.sql.catalyst.util.{toPrettySQL, GeneratedColumn, IdentityColumn, ResolveDefaultColumns, V2ExpressionBuilder} import org.apache.spark.sql.connector.catalog.{Identifier, StagingTableCatalog, SupportsDeleteV2, SupportsNamespaces, SupportsPartitionManagement, SupportsWrite, Table, TableCapability, TableCatalog, TruncatableTable} import org.apache.spark.sql.connector.catalog.index.SupportsIndex import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue} @@ -554,6 +553,9 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat systemScope, pattern) :: Nil + case c: Call => + ExplainOnlySparkPlan(c) :: Nil + case _ => Nil } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExplainOnlySparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExplainOnlySparkPlan.scala new file mode 100644 index 0000000000000..bbf56eaa71184 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExplainOnlySparkPlan.scala @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.v2 + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.trees.LeafLike +import org.apache.spark.sql.execution.SparkPlan + +case class ExplainOnlySparkPlan(toExplain: LogicalPlan) extends SparkPlan with LeafLike[SparkPlan] { + + override def output: Seq[Attribute] = Nil + + override def simpleString(maxFields: Int): String = { + toExplain.simpleString(maxFields) + } + + override protected def doExecute(): RDD[InternalRow] = { + throw new UnsupportedOperationException() + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala index a2539828733fc..0d0258f11efb1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.internal import org.apache.spark.annotation.Unstable import org.apache.spark.sql.{ExperimentalMethods, SparkSession, UDFRegistration, _} import org.apache.spark.sql.artifact.ArtifactManager -import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, FunctionRegistry, ReplaceCharWithVarchar, ResolveSessionCatalog, ResolveTranspose, TableFunctionRegistry} +import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, FunctionRegistry, InvokeProcedures, ReplaceCharWithVarchar, ResolveSessionCatalog, ResolveTranspose, TableFunctionRegistry} import org.apache.spark.sql.catalyst.catalog.{FunctionExpressionBuilder, SessionCatalog} import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.catalyst.optimizer.Optimizer @@ -206,6 +206,7 @@ abstract class BaseSessionStateBuilder( ResolveWriteToStream +: new EvalSubqueriesForTimeTravel +: new ResolveTranspose(session) +: + new InvokeProcedures(session) +: customResolutionRules override val postHocResolutionRules: Seq[Rule[LogicalPlan]] = diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out index 6497a46c68ccd..7c694503056ab 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/keywords.sql.out @@ -32,6 +32,7 @@ BUCKETS false BY false BYTE false CACHE false +CALL true CALLED false CASCADE false CASE true @@ -378,6 +379,7 @@ ANY AS AUTHORIZATION BOTH +CALL CASE CAST CHECK diff --git a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out index 0dfd62599afa6..2c16d961b1313 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out @@ -32,6 +32,7 @@ BUCKETS false BY false BYTE false CACHE false +CALL false CALLED false CASCADE false CASE false diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala new file mode 100644 index 0000000000000..e39a1b7ea340a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala @@ -0,0 +1,654 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import java.util.Collections + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.{SPARK_DOC_ROOT, SparkException, SparkNumberFormatException} +import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId +import org.apache.spark.sql.connector.catalog.{BasicInMemoryTableCatalog, Identifier, InMemoryCatalog} +import org.apache.spark.sql.connector.catalog.procedures.{BoundProcedure, ProcedureParameter, UnboundProcedure} +import org.apache.spark.sql.connector.catalog.procedures.ProcedureParameter.Mode +import org.apache.spark.sql.connector.catalog.procedures.ProcedureParameter.Mode.{IN, INOUT, OUT} +import org.apache.spark.sql.connector.read.{LocalScan, Scan} +import org.apache.spark.sql.errors.DataTypeErrors.{toSQLType, toSQLValue} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DataType, DataTypes, StructField, StructType} +import org.apache.spark.unsafe.types.UTF8String + +class ProcedureSuite extends QueryTest with SharedSparkSession with BeforeAndAfter { + + before { + spark.conf.set(s"spark.sql.catalog.cat", classOf[InMemoryCatalog].getName) + } + + after { + spark.sessionState.catalogManager.reset() + spark.sessionState.conf.unsetConf(s"spark.sql.catalog.cat") + } + + private def catalog: InMemoryCatalog = { + val catalog = spark.sessionState.catalogManager.catalog("cat") + catalog.asInstanceOf[InMemoryCatalog] + } + + test("position arguments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer(sql("CALL cat.ns.sum(5, 5)"), Row(10) :: Nil) + } + + test("named arguments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer(sql("CALL cat.ns.sum(in2 => 3, in1 => 5)"), Row(8) :: Nil) + } + + test("position and named arguments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer(sql("CALL cat.ns.sum(3, in2 => 1)"), Row(4) :: Nil) + } + + test("foldable expressions") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer(sql("CALL cat.ns.sum(1 + 1, in2 => 2)"), Row(4) :: Nil) + checkAnswer(sql("CALL cat.ns.sum(in2 => 1, in1 => 2 + 1)"), Row(4) :: Nil) + checkAnswer(sql("CALL cat.ns.sum((1 + 1) * 2, in2 => (2 + 1) / 3)"), Row(5) :: Nil) + } + + test("type coercion") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundLongSum) + checkAnswer(sql("CALL cat.ns.sum(1, 2)"), Row(3) :: Nil) + checkAnswer(sql("CALL cat.ns.sum(1L, 2)"), Row(3) :: Nil) + checkAnswer(sql("CALL cat.ns.sum(1, 2L)"), Row(3) :: Nil) + } + + test("multiple output rows") { + catalog.createProcedure(Identifier.of(Array("ns"), "complex"), UnboundComplexProcedure) + checkAnswer( + sql("CALL cat.ns.complex('X', 'Y', 3)"), + Row(1, "X1", "Y1") :: Row(2, "X2", "Y2") :: Row(3, "X3", "Y3") :: Nil) + } + + test("parameters with default values") { + catalog.createProcedure(Identifier.of(Array("ns"), "complex"), UnboundComplexProcedure) + checkAnswer(sql("CALL cat.ns.complex()"), Row(1, "A1", "B1") :: Nil) + checkAnswer(sql("CALL cat.ns.complex('X', 'Y')"), Row(1, "X1", "Y1") :: Nil) + } + + test("parameters with invalid default values") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundInvalidDefaultProcedure) + checkError( + exception = intercept[AnalysisException]( + sql("CALL cat.ns.sum()") + ), + condition = "INVALID_DEFAULT_VALUE.DATA_TYPE", + parameters = Map( + "statement" -> "CALL", + "colName" -> toSQLId("in2"), + "defaultValue" -> toSQLValue("B"), + "expectedType" -> toSQLType("INT"), + "actualType" -> toSQLType("STRING"))) + } + + test("IDENTIFIER") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer( + spark.sql("CALL IDENTIFIER(:p1)(1, 2)", Map("p1" -> "cat.ns.sum")), + Row(3) :: Nil) + } + + test("parameterized statements") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkAnswer( + spark.sql("CALL cat.ns.sum(?, ?)", Array(2, 3)), + Row(5) :: Nil) + } + + test("undefined procedure") { + checkError( + exception = intercept[AnalysisException]( + sql("CALL cat.non_exist(1, 2)") + ), + sqlState = Some("38000"), + condition = "FAILED_TO_LOAD_ROUTINE", + parameters = Map("routineName" -> "`cat`.`non_exist`") + ) + } + + test("non-procedure catalog") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[BasicInMemoryTableCatalog].getName) { + checkError( + exception = intercept[AnalysisException]( + sql("CALL testcat.procedure(1, 2)") + ), + condition = "_LEGACY_ERROR_TEMP_1184", + parameters = Map("plugin" -> "testcat", "ability" -> "procedures") + ) + } + } + + test("too many arguments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException]( + sql("CALL cat.ns.sum(1, 2, 3)") + ), + condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + parameters = Map( + "functionName" -> toSQLId("sum"), + "expectedNum" -> "2", + "actualNum" -> "3", + "docroot" -> SPARK_DOC_ROOT)) + } + + test("custom default catalog") { + withSQLConf(SQLConf.DEFAULT_CATALOG.key -> "cat") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + val df = sql("CALL ns.sum(1, 2)") + checkAnswer(df, Row(3) :: Nil) + } + } + + test("custom default catalog and namespace") { + withSQLConf(SQLConf.DEFAULT_CATALOG.key -> "cat") { + catalog.createNamespace(Array("ns"), Collections.emptyMap) + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + sql("USE ns") + val df = sql("CALL sum(1, 2)") + checkAnswer(df, Row(3) :: Nil) + } + } + + test("required parameter not found") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException] { + sql("CALL cat.ns.sum()") + }, + condition = "REQUIRED_PARAMETER_NOT_FOUND", + parameters = Map( + "routineName" -> toSQLId("sum"), + "parameterName" -> toSQLId("in1"), + "index" -> "0")) + } + + test("conflicting position and named parameter assignments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException] { + sql("CALL cat.ns.sum(1, in1 => 2)") + }, + condition = "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.BOTH_POSITIONAL_AND_NAMED", + parameters = Map( + "routineName" -> toSQLId("sum"), + "parameterName" -> toSQLId("in1"))) + } + + test("duplicate named parameter assignments") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException] { + sql("CALL cat.ns.sum(in1 => 1, in1 => 2)") + }, + condition = "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE", + parameters = Map( + "routineName" -> toSQLId("sum"), + "parameterName" -> toSQLId("in1"))) + } + + test("unknown parameter name") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException] { + sql("CALL cat.ns.sum(in1 => 1, in5 => 2)") + }, + condition = "UNRECOGNIZED_PARAMETER_NAME", + parameters = Map( + "routineName" -> toSQLId("sum"), + "argumentName" -> toSQLId("in5"), + "proposal" -> (toSQLId("in1") + " " + toSQLId("in2")))) + } + + test("position parameter after named parameter") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + checkError( + exception = intercept[AnalysisException] { + sql("CALL cat.ns.sum(in1 => 1, 2)") + }, + condition = "UNEXPECTED_POSITIONAL_ARGUMENT", + parameters = Map( + "routineName" -> toSQLId("sum"), + "parameterName" -> toSQLId("in1"))) + } + + test("invalid argument type") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + val call = "CALL cat.ns.sum(1, TIMESTAMP '2016-11-15 20:54:00.000')" + checkError( + exception = intercept[AnalysisException] { + sql(call) + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "CALL", + "paramIndex" -> "second", + "inputSql" -> "\"TIMESTAMP '2016-11-15 20:54:00'\"", + "inputType" -> toSQLType("TIMESTAMP"), + "requiredType" -> toSQLType("INT")), + context = ExpectedContext(fragment = call, start = 0, stop = call.length - 1)) + } + + test("malformed input to implicit cast") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + val call = "CALL cat.ns.sum('A', 2)" + checkError( + exception = intercept[SparkNumberFormatException]( + sql(call) + ), + condition = "CAST_INVALID_INPUT", + parameters = Map( + "expression" -> toSQLValue("A"), + "sourceType" -> toSQLType("STRING"), + "targetType" -> toSQLType("INT")), + context = ExpectedContext(fragment = call, start = 0, stop = call.length - 1)) + } + + test("required parameters after optional") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundInvalidSum) + val e = intercept[SparkException] { + sql("CALL cat.ns.sum(in2 => 1)") + } + assert(e.getMessage.contains("required arguments should come before optional arguments")) + } + + test("INOUT parameters are not supported") { + catalog.createProcedure(Identifier.of(Array("ns"), "procedure"), UnboundInoutProcedure) + val e = intercept[SparkException] { + sql("CALL cat.ns.procedure(1)") + } + assert(e.getMessage.contains(" Unsupported parameter mode: INOUT")) + } + + test("OUT parameters are not supported") { + catalog.createProcedure(Identifier.of(Array("ns"), "procedure"), UnboundOutProcedure) + val e = intercept[SparkException] { + sql("CALL cat.ns.procedure(1)") + } + assert(e.getMessage.contains("Unsupported parameter mode: OUT")) + } + + test("EXPLAIN") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundNonExecutableSum) + val explain1 = sql("EXPLAIN CALL cat.ns.sum(5, 5)").head().get(0) + assert(explain1.toString.contains("cat.ns.sum(5, 5)")) + val explain2 = sql("EXPLAIN EXTENDED CALL cat.ns.sum(10, 10)").head().get(0) + assert(explain2.toString.contains("cat.ns.sum(10, 10)")) + } + + test("void procedure") { + catalog.createProcedure(Identifier.of(Array("ns"), "proc"), UnboundVoidProcedure) + checkAnswer(sql("CALL cat.ns.proc('A', 'B')"), Nil) + } + + test("multi-result procedure") { + catalog.createProcedure(Identifier.of(Array("ns"), "proc"), UnboundMultiResultProcedure) + checkAnswer(sql("CALL cat.ns.proc()"), Row("last") :: Nil) + } + + test("invalid input to struct procedure") { + catalog.createProcedure(Identifier.of(Array("ns"), "proc"), UnboundStructProcedure) + val actualType = + StructType(Seq( + StructField("X", DataTypes.DateType, nullable = false), + StructField("Y", DataTypes.IntegerType, nullable = false))) + val expectedType = StructProcedure.parameters.head.dataType + val call = "CALL cat.ns.proc(named_struct('X', DATE '2011-11-11', 'Y', 2), 'VALUE')" + checkError( + exception = intercept[AnalysisException](sql(call)), + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "CALL", + "paramIndex" -> "first", + "inputSql" -> "\"named_struct(X, DATE '2011-11-11', Y, 2)\"", + "inputType" -> toSQLType(actualType), + "requiredType" -> toSQLType(expectedType)), + context = ExpectedContext(fragment = call, start = 0, stop = call.length - 1)) + } + + test("save execution summary") { + withTable("summary") { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + val result = sql("CALL cat.ns.sum(1, 2)") + result.write.saveAsTable("summary") + checkAnswer(spark.table("summary"), Row(3) :: Nil) + } + } + + object UnboundVoidProcedure extends UnboundProcedure { + override def name: String = "void" + override def description: String = "void procedure" + override def bind(inputType: StructType): BoundProcedure = VoidProcedure + } + + object VoidProcedure extends BoundProcedure { + override def name: String = "void" + + override def description: String = "void procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.StringType).build(), + ProcedureParameter.in("in2", DataTypes.StringType).build() + ) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + Collections.emptyIterator + } + } + + object UnboundMultiResultProcedure extends UnboundProcedure { + override def name: String = "multi" + override def description: String = "multi-result procedure" + override def bind(inputType: StructType): BoundProcedure = MultiResultProcedure + } + + object MultiResultProcedure extends BoundProcedure { + override def name: String = "multi" + + override def description: String = "multi-result procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array() + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + val scans = java.util.Arrays.asList[Scan]( + Result( + new StructType().add("out", DataTypes.IntegerType), + Array(InternalRow(1))), + Result( + new StructType().add("out", DataTypes.StringType), + Array(InternalRow(UTF8String.fromString("last")))) + ) + scans.iterator() + } + } + + object UnboundNonExecutableSum extends UnboundProcedure { + override def name: String = "sum" + override def description: String = "sum integers" + override def bind(inputType: StructType): BoundProcedure = Sum + } + + object NonExecutableSum extends BoundProcedure { + override def name: String = "sum" + + override def description: String = "sum integers" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.IntegerType).build(), + ProcedureParameter.in("in2", DataTypes.IntegerType).build() + ) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + throw new UnsupportedOperationException() + } + } + + object UnboundSum extends UnboundProcedure { + override def name: String = "sum" + override def description: String = "sum integers" + override def bind(inputType: StructType): BoundProcedure = Sum + } + + object Sum extends BoundProcedure { + override def name: String = "sum" + + override def description: String = "sum integers" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.IntegerType).build(), + ProcedureParameter.in("in2", DataTypes.IntegerType).build() + ) + + def outputType: StructType = new StructType().add("out", DataTypes.IntegerType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + val in1 = input.getInt(0) + val in2 = input.getInt(1) + val result = Result(outputType, Array(InternalRow(in1 + in2))) + Collections.singleton[Scan](result).iterator() + } + } + + object UnboundLongSum extends UnboundProcedure { + override def name: String = "long_sum" + override def description: String = "sum longs" + override def bind(inputType: StructType): BoundProcedure = LongSum + } + + object LongSum extends BoundProcedure { + override def name: String = "long_sum" + + override def description: String = "sum longs" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.LongType).build(), + ProcedureParameter.in("in2", DataTypes.LongType).build() + ) + + def outputType: StructType = new StructType().add("out", DataTypes.LongType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + val in1 = input.getLong(0) + val in2 = input.getLong(1) + val result = Result(outputType, Array(InternalRow(in1 + in2))) + Collections.singleton[Scan](result).iterator() + } + } + + object UnboundInvalidSum extends UnboundProcedure { + override def name: String = "invalid" + override def description: String = "sum integers" + override def bind(inputType: StructType): BoundProcedure = InvalidSum + } + + object InvalidSum extends BoundProcedure { + override def name: String = "invalid" + + override def description: String = "sum integers" + + override def isDeterministic: Boolean = false + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.IntegerType).defaultValue("1").build(), + ProcedureParameter.in("in2", DataTypes.IntegerType).build() + ) + + def outputType: StructType = new StructType().add("out", DataTypes.IntegerType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + throw new UnsupportedOperationException() + } + } + + object UnboundInvalidDefaultProcedure extends UnboundProcedure { + override def name: String = "sum" + override def description: String = "invalid default value procedure" + override def bind(inputType: StructType): BoundProcedure = InvalidDefaultProcedure + } + + object InvalidDefaultProcedure extends BoundProcedure { + override def name: String = "sum" + + override def description: String = "invalid default value procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.IntegerType).defaultValue("10").build(), + ProcedureParameter.in("in2", DataTypes.IntegerType).defaultValue("'B'").build() + ) + + def outputType: StructType = new StructType().add("out", DataTypes.IntegerType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + throw new UnsupportedOperationException() + } + } + + object UnboundComplexProcedure extends UnboundProcedure { + override def name: String = "complex" + override def description: String = "complex procedure" + override def bind(inputType: StructType): BoundProcedure = ComplexProcedure + } + + object ComplexProcedure extends BoundProcedure { + override def name: String = "complex" + + override def description: String = "complex procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter.in("in1", DataTypes.StringType).defaultValue("'A'").build(), + ProcedureParameter.in("in2", DataTypes.StringType).defaultValue("'B'").build(), + ProcedureParameter.in("in3", DataTypes.IntegerType).defaultValue("1 + 1 - 1").build() + ) + + def outputType: StructType = new StructType() + .add("out1", DataTypes.IntegerType) + .add("out2", DataTypes.StringType) + .add("out3", DataTypes.StringType) + + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + val in1 = input.getString(0) + val in2 = input.getString(1) + val in3 = input.getInt(2) + + val rows = (1 to in3).map { index => + val v1 = UTF8String.fromString(s"$in1$index") + val v2 = UTF8String.fromString(s"$in2$index") + InternalRow(index, v1, v2) + }.toArray + + val result = Result(outputType, rows) + Collections.singleton[Scan](result).iterator() + } + } + + object UnboundStructProcedure extends UnboundProcedure { + override def name: String = "struct_input" + override def description: String = "struct procedure" + override def bind(inputType: StructType): BoundProcedure = StructProcedure + } + + object StructProcedure extends BoundProcedure { + override def name: String = "struct_input" + + override def description: String = "struct procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + ProcedureParameter + .in( + "in1", + StructType(Seq( + StructField("nested1", DataTypes.IntegerType), + StructField("nested2", DataTypes.StringType)))) + .build(), + ProcedureParameter.in("in2", DataTypes.StringType).build() + ) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + Collections.emptyIterator + } + } + + object UnboundInoutProcedure extends UnboundProcedure { + override def name: String = "procedure" + override def description: String = "inout procedure" + override def bind(inputType: StructType): BoundProcedure = InoutProcedure + } + + object InoutProcedure extends BoundProcedure { + override def name: String = "procedure" + + override def description: String = "inout procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + CustomParameterImpl(INOUT, "in1", DataTypes.IntegerType) + ) + + def outputType: StructType = new StructType().add("out", DataTypes.IntegerType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + throw new UnsupportedOperationException() + } + } + + object UnboundOutProcedure extends UnboundProcedure { + override def name: String = "procedure" + override def description: String = "out procedure" + override def bind(inputType: StructType): BoundProcedure = OutProcedure + } + + object OutProcedure extends BoundProcedure { + override def name: String = "procedure" + + override def description: String = "out procedure" + + override def isDeterministic: Boolean = true + + override def parameters: Array[ProcedureParameter] = Array( + CustomParameterImpl(IN, "in1", DataTypes.IntegerType), + CustomParameterImpl(OUT, "out1", DataTypes.IntegerType) + ) + + def outputType: StructType = new StructType().add("out", DataTypes.IntegerType) + + override def call(input: InternalRow): java.util.Iterator[Scan] = { + throw new UnsupportedOperationException() + } + } + + case class Result(readSchema: StructType, rows: Array[InternalRow]) extends LocalScan + + case class CustomParameterImpl( + mode: Mode, + name: String, + dataType: DataType) extends ProcedureParameter { + override def defaultValueExpression: String = null + override def comment: String = null + } +} diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala index 4bc4116a23da7..dcf3bd8c71731 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala @@ -214,7 +214,7 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer { val sessionHandle = client.openSession(user, "") val infoValue = client.getInfo(sessionHandle, GetInfoType.CLI_ODBC_KEYWORDS) // scalastyle:off line.size.limit - assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") + assert(infoValue.getStringValue == "ADD,AFTER,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,ARCHIVE,ARRAY,AS,ASC,AT,AUTHORIZATION,BEGIN,BETWEEN,BIGINT,BINARY,BINDING,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALL,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CHANGE,CHAR,CHARACTER,CHECK,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONSTRAINT,CONTAINS,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFINED,DEFINER,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,END,ESCAPE,ESCAPED,EVOLUTION,EXCEPT,EXCHANGE,EXCLUDE,EXECUTE,EXISTS,EXPLAIN,EXPORT,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GLOBAL,GRANT,GROUP,GROUPING,HAVING,HOUR,HOURS,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCATION,LOCK,LOCKS,LOGICAL,LONG,MACRO,MAP,MATCHED,MERGE,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NO,NONE,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROPERTIES,PURGE,QUARTER,QUERY,RANGE,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,REDUCE,REFERENCES,REFRESH,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,START,STATISTICS,STORED,STRATIFY,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WEEK,WEEKS,WHEN,WHERE,WHILE,WINDOW,WITH,WITHIN,X,YEAR,YEARS,ZONE") // scalastyle:on line.size.limit } } diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionStateBuilder.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionStateBuilder.scala index 44c1ecd6902ce..dbeb8607facc2 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionStateBuilder.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionStateBuilder.scala @@ -25,7 +25,7 @@ import org.apache.hadoop.hive.ql.exec.{UDAF, UDF} import org.apache.hadoop.hive.ql.udf.generic.{AbstractGenericUDAFResolver, GenericUDF, GenericUDTF} import org.apache.spark.sql._ -import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, ReplaceCharWithVarchar, ResolveSessionCatalog, ResolveTranspose} +import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, InvokeProcedures, ReplaceCharWithVarchar, ResolveSessionCatalog, ResolveTranspose} import org.apache.spark.sql.catalyst.catalog.{ExternalCatalogWithListener, InvalidUDFClassException} import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan @@ -95,6 +95,7 @@ class HiveSessionStateBuilder( new EvalSubqueriesForTimeTravel +: new DetermineTableStats(session) +: new ResolveTranspose(session) +: + new InvokeProcedures(session) +: customResolutionRules override val postHocResolutionRules: Seq[Rule[LogicalPlan]] = From ac34f1de92c6f5cb53d799f00e550a0a204d9eb2 Mon Sep 17 00:00:00 2001 From: Mihailo Milosevic Date: Thu, 19 Sep 2024 11:56:10 +0200 Subject: [PATCH 070/189] [SPARK-48280][SQL][FOLLOW-UP] Add expressions that are built via expressionBuilder to Expression Walker ### What changes were proposed in this pull request? Addition of new expressions to expression walker. This PR also improves descriptions of methods in the Suite. ### Why are the changes needed? It was noticed while debugging that startsWith, endsWith and contains are not tested with this suite and these expressions represent core of collation testing. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Test only. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48162 from mihailom-db/expressionwalkerfollowup. Authored-by: Mihailo Milosevic Signed-off-by: Wenchen Fan --- .../sql/CollationExpressionWalkerSuite.scala | 148 ++++++++++++++---- 1 file changed, 121 insertions(+), 27 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala index 2342722c0bb14..1d23774a51692 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql import java.sql.Timestamp import org.apache.spark.{SparkFunSuite, SparkRuntimeException} +import org.apache.spark.sql.catalyst.analysis.ExpressionBuilder import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.variant.ParseJson import org.apache.spark.sql.internal.SqlApiConf @@ -46,7 +47,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * * @param inputEntry - List of all input entries that need to be generated * @param collationType - Flag defining collation type to use - * @return + * @return - List of data generated for expression instance creation */ def generateData( inputEntry: Seq[Any], @@ -54,23 +55,11 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi inputEntry.map(generateSingleEntry(_, collationType)) } - /** - * Helper function to generate single entry of data as a string. - * @param inputEntry - Single input entry that requires generation - * @param collationType - Flag defining collation type to use - * @return - */ - def generateDataAsStrings( - inputEntry: Seq[AbstractDataType], - collationType: CollationType): Seq[Any] = { - inputEntry.map(generateInputAsString(_, collationType)) - } - /** * Helper function to generate single entry of data. * @param inputEntry - Single input entry that requires generation * @param collationType - Flag defining collation type to use - * @return + * @return - Single input entry data */ def generateSingleEntry( inputEntry: Any, @@ -100,7 +89,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * * @param inputType - Single input literal type that requires generation * @param collationType - Flag defining collation type to use - * @return + * @return - Literal/Expression containing expression ready for evaluation */ def generateLiterals( inputType: AbstractDataType, @@ -116,6 +105,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi } case BooleanType => Literal(true) case _: DatetimeType => Literal(Timestamp.valueOf("2009-07-30 12:58:59")) + case DecimalType => Literal((new Decimal).set(5)) case _: DecimalType => Literal((new Decimal).set(5)) case _: DoubleType => Literal(5.0) case IntegerType | NumericType | IntegralType => Literal(5) @@ -158,11 +148,15 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case MapType => val key = generateLiterals(StringTypeAnyCollation, collationType) val value = generateLiterals(StringTypeAnyCollation, collationType) - Literal.create(Map(key -> value)) + CreateMap(Seq(key, value)) case MapType(keyType, valueType, _) => val key = generateLiterals(keyType, collationType) val value = generateLiterals(valueType, collationType) - Literal.create(Map(key -> value)) + CreateMap(Seq(key, value)) + case AbstractMapType(keyType, valueType) => + val key = generateLiterals(keyType, collationType) + val value = generateLiterals(valueType, collationType) + CreateMap(Seq(key, value)) case StructType => CreateNamedStruct( Seq(Literal("start"), generateLiterals(StringTypeAnyCollation, collationType), @@ -174,7 +168,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * * @param inputType - Single input type that requires generation * @param collationType - Flag defining collation type to use - * @return + * @return - String representation of a input ready for SQL query */ def generateInputAsString( inputType: AbstractDataType, @@ -189,6 +183,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi } case BooleanType => "True" case _: DatetimeType => "date'2016-04-08'" + case DecimalType => "5.0" case _: DecimalType => "5.0" case _: DoubleType => "5.0" case IntegerType | NumericType | IntegralType => "5" @@ -221,6 +216,9 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case MapType(keyType, valueType, _) => "map(" + generateInputAsString(keyType, collationType) + ", " + generateInputAsString(valueType, collationType) + ")" + case AbstractMapType(keyType, valueType) => + "map(" + generateInputAsString(keyType, collationType) + ", " + + generateInputAsString(valueType, collationType) + ")" case StructType => "named_struct( 'start', " + generateInputAsString(StringTypeAnyCollation, collationType) + ", 'end', " + generateInputAsString(StringTypeAnyCollation, collationType) + ")" @@ -234,7 +232,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * * @param inputType - Single input type that requires generation * @param collationType - Flag defining collation type to use - * @return + * @return - String representation for SQL query of a inputType */ def generateInputTypeAsStrings( inputType: AbstractDataType, @@ -244,6 +242,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case BinaryType => "BINARY" case BooleanType => "BOOLEAN" case _: DatetimeType => "DATE" + case DecimalType => "DECIMAL(2, 1)" case _: DecimalType => "DECIMAL(2, 1)" case _: DoubleType => "DOUBLE" case IntegerType | NumericType | IntegralType => "INT" @@ -275,6 +274,9 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case MapType(keyType, valueType, _) => "map<" + generateInputTypeAsStrings(keyType, collationType) + ", " + generateInputTypeAsStrings(valueType, collationType) + ">" + case AbstractMapType(keyType, valueType) => + "map<" + generateInputTypeAsStrings(keyType, collationType) + ", " + + generateInputTypeAsStrings(valueType, collationType) + ">" case StructType => "struct hasStringType(elementType) case TypeCollection(typeCollection) => typeCollection.exists(hasStringType) - case StructType => true case StructType(fields) => fields.exists(sf => hasStringType(sf.dataType)) case _ => false } @@ -310,7 +311,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * Helper function to replace expected parameters with expected input types. * @param inputTypes - Input types generated by ExpectsInputType.inputTypes * @param params - Parameters that are read from expression info - * @return + * @return - List of parameters where Expressions are replaced with input types */ def replaceExpressions(inputTypes: Seq[AbstractDataType], params: Seq[Class[_]]): Seq[Any] = { (inputTypes, params) match { @@ -325,7 +326,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi /** * Helper method to extract relevant expressions that can be walked over. - * @return + * @return - (List of relevant expressions that expect input, List of expressions to skip) */ def extractRelevantExpressions(): (Array[ExpressionInfo], List[String]) = { var expressionCounter = 0 @@ -384,6 +385,47 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi (funInfos, toSkip) } + /** + * Helper method to extract relevant expressions that can be walked over but are built with + * expression builder. + * + * @return - (List of expressions that are relevant builders, List of expressions to skip) + */ + def extractRelevantBuilders(): (Array[ExpressionInfo], List[String]) = { + var builderExpressionCounter = 0 + val funInfos = spark.sessionState.functionRegistry.listFunction().map { funcId => + spark.sessionState.catalog.lookupFunctionInfo(funcId) + }.filter(funInfo => { + // make sure that there is a constructor. + val cl = Utils.classForName(funInfo.getClassName) + cl.isAssignableFrom(classOf[ExpressionBuilder]) + }).filter(funInfo => { + builderExpressionCounter = builderExpressionCounter + 1 + val cl = Utils.classForName(funInfo.getClassName) + val method = cl.getMethod("build", + Utils.classForName("java.lang.String"), + Utils.classForName("scala.collection.Seq")) + var input: Seq[Expression] = Seq.empty + var i = 0 + for (_ <- 1 to 10) { + input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + try { + method.invoke(null, funInfo.getClassName, input).asInstanceOf[ExpectsInputTypes] + } + catch { + case _: Exception => i = i + 1 + } + } + if (i == 10) false + else true + }).toArray + + logInfo("Total number of expression that are built: " + builderExpressionCounter) + logInfo("Number of extracted expressions of relevance: " + funInfos.length) + + (funInfos, List()) + } + /** * Helper function to generate string of an expression suitable for execution. * @param expr - Expression that needs to be converted @@ -441,10 +483,36 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * 5) Otherwise, check if exceptions are the same */ test("SPARK-48280: Expression Walker for expression evaluation") { - val (funInfos, toSkip) = extractRelevantExpressions() + val (funInfosExpr, toSkip) = extractRelevantExpressions() + val (funInfosBuild, _) = extractRelevantBuilders() + val funInfos = funInfosExpr ++ funInfosBuild for (f <- funInfos.filter(f => !toSkip.contains(f.getName))) { - val cl = Utils.classForName(f.getClassName) + val TempCl = Utils.classForName(f.getClassName) + val cl = if (TempCl.isAssignableFrom(classOf[ExpressionBuilder])) { + val clTemp = Utils.classForName(f.getClassName) + val method = clTemp.getMethod("build", + Utils.classForName("java.lang.String"), + Utils.classForName("scala.collection.Seq")) + val instance = { + var input: Seq[Expression] = Seq.empty + var result: Expression = null + for (_ <- 1 to 10) { + input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + try { + val tempResult = method.invoke(null, f.getClassName, input) + if (result == null) result = tempResult.asInstanceOf[Expression] + } + catch { + case _: Exception => + } + } + result + } + instance.getClass + } + else Utils.classForName(f.getClassName) + val headConstructor = cl.getConstructors .zip(cl.getConstructors.map(c => c.getParameters.length)).minBy(a => a._2)._1 val params = headConstructor.getParameters.map(p => p.getType) @@ -526,10 +594,36 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi * 5) Otherwise, check if exceptions are the same */ test("SPARK-48280: Expression Walker for codeGen generation") { - val (funInfos, toSkip) = extractRelevantExpressions() + val (funInfosExpr, toSkip) = extractRelevantExpressions() + val (funInfosBuild, _) = extractRelevantBuilders() + val funInfos = funInfosExpr ++ funInfosBuild for (f <- funInfos.filter(f => !toSkip.contains(f.getName))) { - val cl = Utils.classForName(f.getClassName) + val TempCl = Utils.classForName(f.getClassName) + val cl = if (TempCl.isAssignableFrom(classOf[ExpressionBuilder])) { + val clTemp = Utils.classForName(f.getClassName) + val method = clTemp.getMethod("build", + Utils.classForName("java.lang.String"), + Utils.classForName("scala.collection.Seq")) + val instance = { + var input: Seq[Expression] = Seq.empty + var result: Expression = null + for (_ <- 1 to 10) { + input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + try { + val tempResult = method.invoke(null, f.getClassName, input) + if (result == null) result = tempResult.asInstanceOf[Expression] + } + catch { + case _: Exception => + } + } + result + } + instance.getClass + } + else Utils.classForName(f.getClassName) + val headConstructor = cl.getConstructors .zip(cl.getConstructors.map(c => c.getParameters.length)).minBy(a => a._2)._1 val params = headConstructor.getParameters.map(p => p.getType) From a060c236d314bd2facc73ad26926b59401e5f7aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vladan=20Vasi=C4=87?= Date: Thu, 19 Sep 2024 14:25:53 +0200 Subject: [PATCH 071/189] [SPARK-49667][SQL] Disallowed CS_AI collators with expressions that use StringSearch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? In this PR, I propose to disallow `CS_AI` collated strings in expressions that use `StringsSearch` in their implementation. These expressions are `trim`, `startswith`, `endswith`, `locate`, `instr`, `str_to_map`, `contains`, `replace`, `split_part` and `substring_index`. Currently, these expressions support all possible collations, however, they do not work properly with `CS_AI` collators. This is because there is no support for `CS_AI` search in the ICU's `StringSearch` class which is used to implement these expressions. Therefore, the expressions are not behaving correctly when used with `CS_AI` collators (e.g. currently `startswith('hOtEl' collate unicode_ai, 'Hotel' collate unicode_ai)` returns `true`). ### Why are the changes needed? Proposed changes are necessary in order to achieve correct behavior of the expressions mentioned above. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? This patch was tested by adding a test in the `CollationSuite`. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48121 from vladanvasi-db/vladanvasi-db/cs-ai-collations-expressions-disablement. Authored-by: Vladan Vasić Signed-off-by: Wenchen Fan --- .../sql/catalyst/util/CollationFactory.java | 12 + .../internal/types/AbstractStringType.scala | 9 + .../apache/spark/sql/types/StringType.scala | 3 + .../expressions/complexTypeCreator.scala | 4 +- .../expressions/stringExpressions.scala | 33 +- .../analyzer-results/collations.sql.out | 336 ++++++++++++++++ .../resources/sql-tests/inputs/collations.sql | 14 + .../sql-tests/results/collations.sql.out | 364 ++++++++++++++++++ .../sql/CollationSQLExpressionsSuite.scala | 24 ++ .../sql/CollationStringExpressionsSuite.scala | 251 ++++++++++++ 10 files changed, 1041 insertions(+), 9 deletions(-) diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java index 87558971042e0..d5dbca7eb89bc 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java @@ -921,6 +921,18 @@ public static int collationNameToId(String collationName) throws SparkException return Collation.CollationSpec.collationNameToId(collationName); } + /** + * Returns whether the ICU collation is not Case Sensitive Accent Insensitive + * for the given collation id. + * This method is used in expressions which do not support CS_AI collations. + */ + public static boolean isCaseSensitiveAndAccentInsensitive(int collationId) { + return Collation.CollationSpecICU.fromCollationId(collationId).caseSensitivity == + Collation.CollationSpecICU.CaseSensitivity.CS && + Collation.CollationSpecICU.fromCollationId(collationId).accentSensitivity == + Collation.CollationSpecICU.AccentSensitivity.AI; + } + public static void assertValidProvider(String provider) throws SparkException { if (!SUPPORTED_PROVIDERS.contains(provider.toLowerCase())) { Map params = Map.of( diff --git a/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala b/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala index 05d1701eff74d..dc4ee013fd189 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala @@ -51,3 +51,12 @@ case object StringTypeBinaryLcase extends AbstractStringType { case object StringTypeAnyCollation extends AbstractStringType { override private[sql] def acceptsType(other: DataType): Boolean = other.isInstanceOf[StringType] } + +/** + * Use StringTypeNonCSAICollation for expressions supporting all possible collation types except + * CS_AI collation types. + */ +case object StringTypeNonCSAICollation extends AbstractStringType { + override private[sql] def acceptsType(other: DataType): Boolean = + other.isInstanceOf[StringType] && other.asInstanceOf[StringType].isNonCSAI +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala index eba12c4ff4875..c2dd6cec7ba74 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala @@ -44,6 +44,9 @@ class StringType private (val collationId: Int) extends AtomicType with Serializ private[sql] def supportsLowercaseEquality: Boolean = CollationFactory.fetchCollation(collationId).supportsLowercaseEquality + private[sql] def isNonCSAI: Boolean = + !CollationFactory.isCaseSensitiveAndAccentInsensitive(collationId) + private[sql] def isUTF8BinaryCollation: Boolean = collationId == CollationFactory.UTF8_BINARY_COLLATION_ID diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala index ba1beab28d9a7..b8b47f2763f5b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern._ import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeNonCSAICollation import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.ArrayImplicits._ @@ -579,7 +579,7 @@ case class StringToMap(text: Expression, pairDelim: Expression, keyValueDelim: E override def third: Expression = keyValueDelim override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) override def dataType: DataType = MapType(first.dataType, first.dataType) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index e75df87994f0e..da6d786efb4e3 100755 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern, UPPER_OR_LO import org.apache.spark.sql.catalyst.util.{ArrayData, CharsetProvider, CollationFactory, CollationSupport, GenericArrayData, TypeUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeAnyCollation} +import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeAnyCollation, StringTypeNonCSAICollation} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.UTF8StringBuilder import org.apache.spark.unsafe.array.ByteArrayMethods @@ -609,6 +609,8 @@ case class Contains(left: Expression, right: Expression) extends StringPredicate defineCodeGen(ctx, ev, (c1, c2) => CollationSupport.Contains.genCode(c1, c2, collationId)) } + override def inputTypes : Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation) override protected def withNewChildrenInternal( newLeft: Expression, newRight: Expression): Contains = copy(left = newLeft, right = newRight) } @@ -650,6 +652,10 @@ case class StartsWith(left: Expression, right: Expression) extends StringPredica defineCodeGen(ctx, ev, (c1, c2) => CollationSupport.StartsWith.genCode(c1, c2, collationId)) } + + override def inputTypes : Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) + override protected def withNewChildrenInternal( newLeft: Expression, newRight: Expression): StartsWith = copy(left = newLeft, right = newRight) } @@ -691,6 +697,10 @@ case class EndsWith(left: Expression, right: Expression) extends StringPredicate defineCodeGen(ctx, ev, (c1, c2) => CollationSupport.EndsWith.genCode(c1, c2, collationId)) } + + override def inputTypes : Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) + override protected def withNewChildrenInternal( newLeft: Expression, newRight: Expression): EndsWith = copy(left = newLeft, right = newRight) } @@ -919,7 +929,7 @@ case class StringReplace(srcExpr: Expression, searchExpr: Expression, replaceExp override def dataType: DataType = srcExpr.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) override def first: Expression = srcExpr override def second: Expression = searchExpr override def third: Expression = replaceExpr @@ -1167,7 +1177,7 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac override def dataType: DataType = srcExpr.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) override def first: Expression = srcExpr override def second: Expression = matchingExpr override def third: Expression = replaceExpr @@ -1394,6 +1404,9 @@ case class StringTrim(srcStr: Expression, trimStr: Option[Expression] = None) override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrim.exec(srcString, trimString, collationId) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation) + override protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = copy( srcStr = newChildren.head, @@ -1501,6 +1514,9 @@ case class StringTrimLeft(srcStr: Expression, trimStr: Option[Expression] = None override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrimLeft.exec(srcString, trimString, collationId) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation) + override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): StringTrimLeft = copy( @@ -1561,6 +1577,9 @@ case class StringTrimRight(srcStr: Expression, trimStr: Option[Expression] = Non override def doEval(srcString: UTF8String, trimString: UTF8String): UTF8String = CollationSupport.StringTrimRight.exec(srcString, trimString, collationId) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation) + override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): StringTrimRight = copy( @@ -1595,7 +1614,7 @@ case class StringInstr(str: Expression, substr: Expression) override def right: Expression = substr override def dataType: DataType = IntegerType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation) override def nullSafeEval(string: Any, sub: Any): Any = { CollationSupport.StringInstr. @@ -1643,7 +1662,7 @@ case class SubstringIndex(strExpr: Expression, delimExpr: Expression, countExpr: override def dataType: DataType = strExpr.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, IntegerType) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, IntegerType) override def first: Expression = strExpr override def second: Expression = delimExpr override def third: Expression = countExpr @@ -1701,7 +1720,7 @@ case class StringLocate(substr: Expression, str: Expression, start: Expression) override def nullable: Boolean = substr.nullable || str.nullable override def dataType: DataType = IntegerType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, IntegerType) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, IntegerType) override def eval(input: InternalRow): Any = { val s = start.eval(input) @@ -3463,7 +3482,7 @@ case class SplitPart ( false) override def nodeName: String = "split_part" override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, IntegerType) + Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, IntegerType) def children: Seq[Expression] = Seq(str, delimiter, partNum) protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = { copy(str = newChildren.apply(0), delimiter = newChildren.apply(1), diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out index 83c9ebfef4b25..eed7fa73ab698 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/collations.sql.out @@ -436,6 +436,30 @@ Project [str_to_map(collate(text#x, utf8_binary), collate(pairDelim#x, utf8_bina +- Relation spark_catalog.default.t4[text#x,pairDelim#x,keyValueDelim#x] parquet +-- !query +select str_to_map(text collate unicode_ai, pairDelim collate unicode_ai, keyValueDelim collate unicode_ai) from t4 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(text, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"str_to_map(collate(text, unicode_ai), collate(pairDelim, unicode_ai), collate(keyValueDelim, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 106, + "fragment" : "str_to_map(text collate unicode_ai, pairDelim collate unicode_ai, keyValueDelim collate unicode_ai)" + } ] +} + + -- !query drop table t4 -- !query analysis @@ -820,6 +844,30 @@ Project [split_part(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, ut +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select split_part(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"split_part(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 2)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 83, + "fragment" : "split_part(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2)" + } ] +} + + -- !query select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5 -- !query analysis @@ -883,6 +931,30 @@ Project [Contains(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8 +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select contains(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"contains(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 78, + "fragment" : "contains(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5 -- !query analysis @@ -946,6 +1018,30 @@ Project [substring_index(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase# +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select substring_index(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"substring_index(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 2)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 88, + "fragment" : "substring_index(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2)" + } ] +} + + -- !query select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5 -- !query analysis @@ -1009,6 +1105,30 @@ Project [instr(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lc +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select instr(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"instr(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "instr(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5 -- !query analysis @@ -1135,6 +1255,30 @@ Project [StartsWith(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, ut +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select startswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"startswith(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 80, + "fragment" : "startswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5 -- !query analysis @@ -1190,6 +1334,30 @@ Project [translate(cast(utf8_binary#x as string collate UTF8_LCASE), collate(SQL +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select translate(utf8_binary, 'SQL' collate unicode_ai, '12345' collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"utf8_binary\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"translate(utf8_binary, collate(SQL, unicode_ai), collate(12345, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 83, + "fragment" : "translate(utf8_binary, 'SQL' collate unicode_ai, '12345' collate unicode_ai)" + } ] +} + + -- !query select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5 -- !query analysis @@ -1253,6 +1421,30 @@ Project [replace(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_ +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select replace(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 'abc') from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"replace(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), abc)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 84, + "fragment" : "replace(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 'abc')" + } ] +} + + -- !query select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5 -- !query analysis @@ -1316,6 +1508,30 @@ Project [EndsWith(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8 +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select endswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"endswith(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 78, + "fragment" : "endswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5 -- !query analysis @@ -2039,6 +2255,30 @@ Project [locate(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_l +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select locate(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 3) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"locate(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 3)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 79, + "fragment" : "locate(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 3)" + } ] +} + + -- !query select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5 -- !query analysis @@ -2102,6 +2342,30 @@ Project [trim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, utf +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select TRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(BOTH collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 74, + "fragment" : "TRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5 -- !query analysis @@ -2165,6 +2429,30 @@ Project [btrim(collate(utf8_binary#x, utf8_lcase), collate(utf8_lcase#x, utf8_lc +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select BTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(BOTH collate(utf8_lcase, unicode_ai) FROM collate(utf8_binary, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "BTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5 -- !query analysis @@ -2228,6 +2516,30 @@ Project [ltrim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, ut +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select LTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(LEADING collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "LTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5 -- !query analysis @@ -2291,6 +2603,30 @@ Project [rtrim(collate(utf8_lcase#x, utf8_lcase), Some(collate(utf8_binary#x, ut +- Relation spark_catalog.default.t5[s#x,utf8_binary#x,utf8_lcase#x] parquet +-- !query +select RTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(TRAILING collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "RTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5 -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/collations.sql b/sql/core/src/test/resources/sql-tests/inputs/collations.sql index 183577b83971b..f3a42fd3e1f12 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/collations.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/collations.sql @@ -99,6 +99,7 @@ insert into t4 values('a:1,b:2,c:3', ',', ':'); select str_to_map(text, pairDelim, keyValueDelim) from t4; select str_to_map(text collate utf8_binary, pairDelim collate utf8_lcase, keyValueDelim collate utf8_binary) from t4; select str_to_map(text collate utf8_binary, pairDelim collate utf8_binary, keyValueDelim collate utf8_binary) from t4; +select str_to_map(text collate unicode_ai, pairDelim collate unicode_ai, keyValueDelim collate unicode_ai) from t4; drop table t4; @@ -159,6 +160,7 @@ select split_part(s, utf8_binary, 1) from t5; select split_part(utf8_binary collate utf8_binary, s collate utf8_lcase, 1) from t5; select split_part(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5; select split_part(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5; +select split_part(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2) from t5; select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5; select split_part(utf8_binary, 'a' collate utf8_lcase, 3), split_part(utf8_lcase, 'a' collate utf8_binary, 3) from t5; @@ -168,6 +170,7 @@ select contains(s, utf8_binary) from t5; select contains(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select contains(utf8_binary, utf8_lcase collate utf8_binary) from t5; select contains(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select contains(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5; select contains(utf8_binary, 'AaAA' collate utf8_lcase), contains(utf8_lcase, 'AAa' collate utf8_binary) from t5; @@ -177,6 +180,7 @@ select substring_index(s, utf8_binary,1) from t5; select substring_index(utf8_binary collate utf8_binary, s collate utf8_lcase, 3) from t5; select substring_index(utf8_binary, utf8_lcase collate utf8_binary, 2) from t5; select substring_index(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 2) from t5; +select substring_index(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2) from t5; select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5; select substring_index(utf8_binary, 'AaAA' collate utf8_lcase, 2), substring_index(utf8_lcase, 'AAa' collate utf8_binary, 2) from t5; @@ -186,6 +190,7 @@ select instr(s, utf8_binary) from t5; select instr(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select instr(utf8_binary, utf8_lcase collate utf8_binary) from t5; select instr(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select instr(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5; select instr(utf8_binary, 'AaAA' collate utf8_lcase), instr(utf8_lcase, 'AAa' collate utf8_binary) from t5; @@ -204,6 +209,7 @@ select startswith(s, utf8_binary) from t5; select startswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select startswith(utf8_binary, utf8_lcase collate utf8_binary) from t5; select startswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select startswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5; select startswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), startswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5; @@ -212,6 +218,7 @@ select translate(utf8_lcase, utf8_lcase, '12345') from t5; select translate(utf8_binary, utf8_lcase, '12345') from t5; select translate(utf8_binary, 'aBc' collate utf8_lcase, '12345' collate utf8_binary) from t5; select translate(utf8_binary, 'SQL' collate utf8_lcase, '12345' collate utf8_lcase) from t5; +select translate(utf8_binary, 'SQL' collate unicode_ai, '12345' collate unicode_ai) from t5; select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5; select translate(utf8_lcase, 'aBc' collate utf8_binary, '12345'), translate(utf8_binary, 'aBc' collate utf8_lcase, '12345') from t5; @@ -221,6 +228,7 @@ select replace(s, utf8_binary, 'abc') from t5; select replace(utf8_binary collate utf8_binary, s collate utf8_lcase, 'abc') from t5; select replace(utf8_binary, utf8_lcase collate utf8_binary, 'abc') from t5; select replace(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 'abc') from t5; +select replace(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 'abc') from t5; select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5; select replace(utf8_binary, 'aaAaaAaA' collate utf8_lcase, 'abc'), replace(utf8_lcase, 'aaAaaAaA' collate utf8_binary, 'abc') from t5; @@ -230,6 +238,7 @@ select endswith(s, utf8_binary) from t5; select endswith(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select endswith(utf8_binary, utf8_lcase collate utf8_binary) from t5; select endswith(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select endswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5; select endswith(utf8_binary, 'aaAaaAaA' collate utf8_lcase), endswith(utf8_lcase, 'aaAaaAaA' collate utf8_binary) from t5; @@ -364,6 +373,7 @@ select locate(s, utf8_binary) from t5; select locate(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select locate(utf8_binary, utf8_lcase collate utf8_binary) from t5; select locate(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase, 3) from t5; +select locate(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 3) from t5; select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5; select locate(utf8_binary, 'AaAA' collate utf8_lcase, 4), locate(utf8_lcase, 'AAa' collate utf8_binary, 4) from t5; @@ -373,6 +383,7 @@ select TRIM(s, utf8_binary) from t5; select TRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select TRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; select TRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select TRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5; select TRIM('ABc' collate utf8_lcase, utf8_binary), TRIM('AAa' collate utf8_binary, utf8_lcase) from t5; -- StringTrimBoth @@ -381,6 +392,7 @@ select BTRIM(s, utf8_binary) from t5; select BTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select BTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; select BTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select BTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5; select BTRIM('ABc' collate utf8_lcase, utf8_binary), BTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; -- StringTrimLeft @@ -389,6 +401,7 @@ select LTRIM(s, utf8_binary) from t5; select LTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select LTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; select LTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select LTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5; select LTRIM('ABc' collate utf8_lcase, utf8_binary), LTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; -- StringTrimRight @@ -397,6 +410,7 @@ select RTRIM(s, utf8_binary) from t5; select RTRIM(utf8_binary collate utf8_binary, s collate utf8_lcase) from t5; select RTRIM(utf8_binary, utf8_lcase collate utf8_binary) from t5; select RTRIM(utf8_binary collate utf8_lcase, utf8_lcase collate utf8_lcase) from t5; +select RTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5; select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5; select RTRIM('ABc' collate utf8_lcase, utf8_binary), RTRIM('AAa' collate utf8_binary, utf8_lcase) from t5; diff --git a/sql/core/src/test/resources/sql-tests/results/collations.sql.out b/sql/core/src/test/resources/sql-tests/results/collations.sql.out index ea5564aafe96f..5999bf20f6884 100644 --- a/sql/core/src/test/resources/sql-tests/results/collations.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/collations.sql.out @@ -480,6 +480,32 @@ struct +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(text, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"str_to_map(collate(text, unicode_ai), collate(pairDelim, unicode_ai), collate(keyValueDelim, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 106, + "fragment" : "str_to_map(text collate unicode_ai, pairDelim collate unicode_ai, keyValueDelim collate unicode_ai)" + } ] +} + + -- !query drop table t4 -- !query schema @@ -1021,6 +1047,32 @@ struct +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"split_part(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 2)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 83, + "fragment" : "split_part(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2)" + } ] +} + + -- !query select split_part(utf8_binary, 'a', 3), split_part(utf8_lcase, 'a', 3) from t5 -- !query schema @@ -1148,6 +1200,32 @@ true true +-- !query +select contains(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"contains(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 78, + "fragment" : "contains(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select contains(utf8_binary, 'a'), contains(utf8_lcase, 'a') from t5 -- !query schema @@ -1275,6 +1353,32 @@ kitten İo +-- !query +select substring_index(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"substring_index(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 2)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 88, + "fragment" : "substring_index(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 2)" + } ] +} + + -- !query select substring_index(utf8_binary, 'a', 2), substring_index(utf8_lcase, 'a', 2) from t5 -- !query schema @@ -1402,6 +1506,32 @@ struct +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"instr(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "instr(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select instr(utf8_binary, 'a'), instr(utf8_lcase, 'a') from t5 -- !query schema @@ -1656,6 +1786,32 @@ true true +-- !query +select startswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"startswith(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 80, + "fragment" : "startswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select startswith(utf8_binary, 'aaAaaAaA'), startswith(utf8_lcase, 'aaAaaAaA') from t5 -- !query schema @@ -1763,6 +1919,32 @@ kitten İo +-- !query +select translate(utf8_binary, 'SQL' collate unicode_ai, '12345' collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"utf8_binary\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"translate(utf8_binary, collate(SQL, unicode_ai), collate(12345, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 83, + "fragment" : "translate(utf8_binary, 'SQL' collate unicode_ai, '12345' collate unicode_ai)" + } ] +} + + -- !query select translate(utf8_lcase, 'aaAaaAaA', '12345'), translate(utf8_binary, 'aaAaaAaA', '12345') from t5 -- !query schema @@ -1890,6 +2072,32 @@ bbabcbabcabcbabc kitten +-- !query +select replace(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 'abc') from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"replace(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), abc)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 84, + "fragment" : "replace(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 'abc')" + } ] +} + + -- !query select replace(utf8_binary, 'aaAaaAaA', 'abc'), replace(utf8_lcase, 'aaAaaAaA', 'abc') from t5 -- !query schema @@ -2017,6 +2225,32 @@ true true +-- !query +select endswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"endswith(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 78, + "fragment" : "endswith(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select endswith(utf8_binary, 'aaAaaAaA'), endswith(utf8_lcase, 'aaAaaAaA') from t5 -- !query schema @@ -3570,6 +3804,32 @@ struct +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"locate(collate(utf8_binary, unicode_ai), collate(utf8_lcase, unicode_ai), 3)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 79, + "fragment" : "locate(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai, 3)" + } ] +} + + -- !query select locate(utf8_binary, 'a'), locate(utf8_lcase, 'a') from t5 -- !query schema @@ -3685,6 +3945,32 @@ QL sitTing +-- !query +select TRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(BOTH collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 74, + "fragment" : "TRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select TRIM('ABc', utf8_binary), TRIM('ABc', utf8_lcase) from t5 -- !query schema @@ -3812,6 +4098,32 @@ park İ +-- !query +select BTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_binary, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(BOTH collate(utf8_lcase, unicode_ai) FROM collate(utf8_binary, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "BTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select BTRIM('ABc', utf8_binary), BTRIM('ABc', utf8_lcase) from t5 -- !query schema @@ -3927,6 +4239,32 @@ QL sitTing +-- !query +select LTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(LEADING collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "LTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select LTRIM('ABc', utf8_binary), LTRIM('ABc', utf8_lcase) from t5 -- !query schema @@ -4042,6 +4380,32 @@ SQL sitTing +-- !query +select RTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai) from t5 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"collate(utf8_lcase, unicode_ai)\"", + "inputType" : "\"STRING COLLATE UNICODE_AI\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"TRIM(TRAILING collate(utf8_binary, unicode_ai) FROM collate(utf8_lcase, unicode_ai))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 75, + "fragment" : "RTRIM(utf8_binary collate unicode_ai, utf8_lcase collate unicode_ai)" + } ] +} + + -- !query select RTRIM('ABc', utf8_binary), RTRIM('ABc', utf8_lcase) from t5 -- !query schema diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala index f8cd840ecdbb9..941d5cd31db40 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala @@ -982,6 +982,7 @@ class CollationSQLExpressionsSuite StringToMapTestCase("1/AX2/BX3/C", "x", "/", "UNICODE_CI", Map("1" -> "A", "2" -> "B", "3" -> "C")) ) + val unsupportedTestCase = StringToMapTestCase("a:1,b:2,c:3", "?", "?", "UNICODE_AI", null) testCases.foreach(t => { // Unit test. val text = Literal.create(t.text, StringType(t.collation)) @@ -996,6 +997,29 @@ class CollationSQLExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(dataType)) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select str_to_map('${unsupportedTestCase.text}', '${unsupportedTestCase.pairDelim}', " + + s"'${unsupportedTestCase.keyValueDelim}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> ("\"str_to_map('a:1,b:2,c:3' collate UNICODE_AI, " + + "'?' collate UNICODE_AI, '?' collate UNICODE_AI)\""), + "paramIndex" -> "first", + "inputSql" -> "\"'a:1,b:2,c:3' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext( + fragment = "str_to_map('a:1,b:2,c:3', '?', '?')", + start = 7, + stop = 41)) + } } test("Support RaiseError misc expression with collation") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala index 6804411d470b9..fe9872ddaf575 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationStringExpressionsSuite.scala @@ -98,6 +98,7 @@ class CollationStringExpressionsSuite SplitPartTestCase("1a2", "A", 2, "UTF8_LCASE", "2"), SplitPartTestCase("1a2", "A", 2, "UNICODE_CI", "2") ) + val unsupportedTestCase = SplitPartTestCase("1a2", "a", 2, "UNICODE_AI", "2") testCases.foreach(t => { // Unit test. val str = Literal.create(t.str, StringType(t.collation)) @@ -111,6 +112,26 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select split_part('${unsupportedTestCase.str}', '${unsupportedTestCase.delimiter}', " + + s"${unsupportedTestCase.partNum})" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"split_part('1a2' collate UNICODE_AI, 'a' collate UNICODE_AI, 2)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'1a2' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "split_part('1a2', 'a', 2)", start = 7, stop = 31) + ) + } } test("Support `StringSplitSQL` string expression with collation") { @@ -166,6 +187,7 @@ class CollationStringExpressionsSuite ContainsTestCase("abcde", "FGH", "UTF8_LCASE", false), ContainsTestCase("abcde", "BCD", "UNICODE_CI", true) ) + val unsupportedTestCase = ContainsTestCase("abcde", "A", "UNICODE_AI", false) testCases.foreach(t => { // Unit test. val left = Literal.create(t.left, StringType(t.collation)) @@ -178,6 +200,25 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(BooleanType)) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select contains('${unsupportedTestCase.left}', '${unsupportedTestCase.right}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"contains('abcde' collate UNICODE_AI, 'A' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'abcde' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "contains('abcde', 'A')", start = 7, stop = 28) + ) + } } test("Support `SubstringIndex` expression with collation") { @@ -194,6 +235,7 @@ class CollationStringExpressionsSuite SubstringIndexTestCase("aaaaaaaaaa", "aa", 2, "UNICODE", "a"), SubstringIndexTestCase("wwwmapacheMorg", "M", -2, "UNICODE_CI", "apacheMorg") ) + val unsupportedTestCase = SubstringIndexTestCase("abacde", "a", 2, "UNICODE_AI", "cde") testCases.foreach(t => { // Unit test. val strExpr = Literal.create(t.strExpr, StringType(t.collation)) @@ -207,6 +249,29 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select substring_index('${unsupportedTestCase.strExpr}', " + + s"'${unsupportedTestCase.delimExpr}', ${unsupportedTestCase.countExpr})" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> ("\"substring_index('abacde' collate UNICODE_AI, " + + "'a' collate UNICODE_AI, 2)\""), + "paramIndex" -> "first", + "inputSql" -> "\"'abacde' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext( + fragment = "substring_index('abacde', 'a', 2)", + start = 7, + stop = 39)) + } } test("Support `StringInStr` string expression with collation") { @@ -219,6 +284,7 @@ class CollationStringExpressionsSuite StringInStrTestCase("test大千世界X大千世界", "界x", "UNICODE_CI", 8), StringInStrTestCase("abİo12", "i̇o", "UNICODE_CI", 3) ) + val unsupportedTestCase = StringInStrTestCase("a", "abcde", "UNICODE_AI", 0) testCases.foreach(t => { // Unit test. val str = Literal.create(t.str, StringType(t.collation)) @@ -231,6 +297,25 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(IntegerType)) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select instr('${unsupportedTestCase.str}', '${unsupportedTestCase.substr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"instr('a' collate UNICODE_AI, 'abcde' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'a' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "instr('a', 'abcde')", start = 7, stop = 25) + ) + } } test("Support `FindInSet` string expression with collation") { @@ -264,6 +349,7 @@ class CollationStringExpressionsSuite StartsWithTestCase("abcde", "FGH", "UTF8_LCASE", false), StartsWithTestCase("abcde", "ABC", "UNICODE_CI", true) ) + val unsupportedTestCase = StartsWithTestCase("abcde", "A", "UNICODE_AI", false) testCases.foreach(t => { // Unit test. val left = Literal.create(t.left, StringType(t.collation)) @@ -276,6 +362,25 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(BooleanType)) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select startswith('${unsupportedTestCase.left}', '${unsupportedTestCase.right}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"startswith('abcde' collate UNICODE_AI, 'A' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'abcde' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "startswith('abcde', 'A')", start = 7, stop = 30) + ) + } } test("Support `StringTranslate` string expression with collation") { @@ -291,6 +396,7 @@ class CollationStringExpressionsSuite StringTranslateTestCase("Translate", "Rn", "\u0000\u0000", "UNICODE", "Traslate"), StringTranslateTestCase("Translate", "Rn", "1234", "UNICODE_CI", "T1a2slate") ) + val unsupportedTestCase = StringTranslateTestCase("ABC", "AB", "12", "UNICODE_AI", "12C") testCases.foreach(t => { // Unit test. val srcExpr = Literal.create(t.srcExpr, StringType(t.collation)) @@ -304,6 +410,27 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select translate('${unsupportedTestCase.srcExpr}', " + + s"'${unsupportedTestCase.matchingExpr}', '${unsupportedTestCase.replaceExpr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> ("\"translate('ABC' collate UNICODE_AI, 'AB' collate UNICODE_AI, " + + "'12' collate UNICODE_AI)\""), + "paramIndex" -> "first", + "inputSql" -> "\"'ABC' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "translate('ABC', 'AB', '12')", start = 7, stop = 34) + ) + } } test("Support `StringReplace` string expression with collation") { @@ -321,6 +448,7 @@ class CollationStringExpressionsSuite StringReplaceTestCase("abi̇o12i̇o", "İo", "yy", "UNICODE_CI", "abyy12yy"), StringReplaceTestCase("abİo12i̇o", "i̇o", "xx", "UNICODE_CI", "abxx12xx") ) + val unsupportedTestCase = StringReplaceTestCase("abcde", "A", "B", "UNICODE_AI", "abcde") testCases.foreach(t => { // Unit test. val srcExpr = Literal.create(t.srcExpr, StringType(t.collation)) @@ -334,6 +462,27 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select replace('${unsupportedTestCase.srcExpr}', '${unsupportedTestCase.searchExpr}', " + + s"'${unsupportedTestCase.replaceExpr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> ("\"replace('abcde' collate UNICODE_AI, 'A' collate UNICODE_AI, " + + "'B' collate UNICODE_AI)\""), + "paramIndex" -> "first", + "inputSql" -> "\"'abcde' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "replace('abcde', 'A', 'B')", start = 7, stop = 32) + ) + } } test("Support `EndsWith` string expression with collation") { @@ -344,6 +493,7 @@ class CollationStringExpressionsSuite EndsWithTestCase("abcde", "FGH", "UTF8_LCASE", false), EndsWithTestCase("abcde", "CDE", "UNICODE_CI", true) ) + val unsupportedTestCase = EndsWithTestCase("abcde", "A", "UNICODE_AI", false) testCases.foreach(t => { // Unit test. val left = Literal.create(t.left, StringType(t.collation)) @@ -355,6 +505,25 @@ class CollationStringExpressionsSuite checkAnswer(sql(query), Row(t.result)) assert(sql(query).schema.fields.head.dataType.sameType(BooleanType)) } + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select endswith('${unsupportedTestCase.left}', '${unsupportedTestCase.right}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"endswith('abcde' collate UNICODE_AI, 'A' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'abcde' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "endswith('abcde', 'A')", start = 7, stop = 28) + ) + } }) } @@ -1097,6 +1266,7 @@ class CollationStringExpressionsSuite StringLocateTestCase("aa", "Aaads", 0, "UNICODE_CI", 0), StringLocateTestCase("界x", "test大千世界X大千世界", 1, "UNICODE_CI", 8) ) + val unsupportedTestCase = StringLocateTestCase("aa", "Aaads", 0, "UNICODE_AI", 1) testCases.foreach(t => { // Unit test. val substr = Literal.create(t.substr, StringType(t.collation)) @@ -1110,6 +1280,26 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(IntegerType)) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val query = + s"select locate('${unsupportedTestCase.substr}', '${unsupportedTestCase.str}', " + + s"${unsupportedTestCase.start})" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"locate('aa' collate UNICODE_AI, 'Aaads' collate UNICODE_AI, 0)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'aa' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "locate('aa', 'Aaads', 0)", start = 7, stop = 30) + ) + } } test("Support `StringTrimLeft` string expression with collation") { @@ -1124,6 +1314,7 @@ class CollationStringExpressionsSuite StringTrimLeftTestCase("xxasdxx", Some("y"), "UNICODE", "xxasdxx"), StringTrimLeftTestCase(" asd ", None, "UNICODE_CI", "asd ") ) + val unsupportedTestCase = StringTrimLeftTestCase("xxasdxx", Some("x"), "UNICODE_AI", null) testCases.foreach(t => { // Unit test. val srcStr = Literal.create(t.srcStr, StringType(t.collation)) @@ -1137,6 +1328,25 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val trimString = s"'${unsupportedTestCase.trimStr.get}', " + val query = s"select ltrim($trimString'${unsupportedTestCase.srcStr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"TRIM(LEADING 'x' collate UNICODE_AI FROM 'xxasdxx' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'xxasdxx' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "ltrim('x', 'xxasdxx')", start = 7, stop = 27) + ) + } } test("Support `StringTrimRight` string expression with collation") { @@ -1151,6 +1361,7 @@ class CollationStringExpressionsSuite StringTrimRightTestCase("xxasdxx", Some("y"), "UNICODE", "xxasdxx"), StringTrimRightTestCase(" asd ", None, "UNICODE_CI", " asd") ) + val unsupportedTestCase = StringTrimRightTestCase("xxasdxx", Some("x"), "UNICODE_AI", "xxasd") testCases.foreach(t => { // Unit test. val srcStr = Literal.create(t.srcStr, StringType(t.collation)) @@ -1164,6 +1375,26 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val trimString = s"'${unsupportedTestCase.trimStr.get}', " + val query = s"select rtrim($trimString'${unsupportedTestCase.srcStr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> ("\"TRIM(TRAILING 'x' collate UNICODE_AI FROM 'xxasdxx'" + + " collate UNICODE_AI)\""), + "paramIndex" -> "first", + "inputSql" -> "\"'xxasdxx' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "rtrim('x', 'xxasdxx')", start = 7, stop = 27) + ) + } } test("Support `StringTrim` string expression with collation") { @@ -1178,6 +1409,7 @@ class CollationStringExpressionsSuite StringTrimTestCase("xxasdxx", Some("y"), "UNICODE", "xxasdxx"), StringTrimTestCase(" asd ", None, "UNICODE_CI", "asd") ) + val unsupportedTestCase = StringTrimTestCase("xxasdxx", Some("x"), "UNICODE_AI", "asd") testCases.foreach(t => { // Unit test. val srcStr = Literal.create(t.srcStr, StringType(t.collation)) @@ -1191,6 +1423,25 @@ class CollationStringExpressionsSuite assert(sql(query).schema.fields.head.dataType.sameType(StringType(t.collation))) } }) + // Test unsupported collation. + withSQLConf(SQLConf.DEFAULT_COLLATION.key -> unsupportedTestCase.collation) { + val trimString = s"'${unsupportedTestCase.trimStr.get}', " + val query = s"select trim($trimString'${unsupportedTestCase.srcStr}')" + checkError( + exception = intercept[AnalysisException] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + sqlState = Some("42K09"), + parameters = Map( + "sqlExpr" -> "\"TRIM(BOTH 'x' collate UNICODE_AI FROM 'xxasdxx' collate UNICODE_AI)\"", + "paramIndex" -> "first", + "inputSql" -> "\"'xxasdxx' collate UNICODE_AI\"", + "inputType" -> "\"STRING COLLATE UNICODE_AI\"", + "requiredType" -> "\"STRING\""), + context = ExpectedContext(fragment = "trim('x', 'xxasdxx')", start = 7, stop = 26) + ) + } } test("Support `StringTrimBoth` string expression with collation") { From 4068fbcc0de59154db9bdeb1296bd24059db9f42 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 21:00:57 +0800 Subject: [PATCH 072/189] [SPARK-49717][SQL][TESTS] Function parity test ignore private[xxx] functions ### What changes were proposed in this pull request? Function parity test ignore private functions ### Why are the changes needed? existing test is based on `java.lang.reflect.Modifier` which cannot properly handle `private[xxx]` ### Does this PR introduce _any_ user-facing change? no, test only ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48163 from zhengruifeng/df_func_test. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- .../apache/spark/sql/DataFrameFunctionsSuite.scala | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index f16171940df21..0842b92e5d53c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala @@ -17,10 +17,10 @@ package org.apache.spark.sql -import java.lang.reflect.Modifier import java.nio.charset.StandardCharsets import java.sql.{Date, Timestamp} +import scala.reflect.runtime.universe.runtimeMirror import scala.util.Random import org.apache.spark.{QueryContextType, SPARK_DOC_ROOT, SparkException, SparkRuntimeException} @@ -82,7 +82,6 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession { "bucket", "days", "hours", "months", "years", // Datasource v2 partition transformations "product", // Discussed in https://github.com/apache/spark/pull/30745 "unwrap_udt", - "collect_top_k", "timestamp_add", "timestamp_diff" ) @@ -92,10 +91,13 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession { val word_pattern = """\w*""" // Set of DataFrame functions in org.apache.spark.sql.functions - val dataFrameFunctions = functions.getClass - .getDeclaredMethods - .filter(m => Modifier.isPublic(m.getModifiers)) - .map(_.getName) + val dataFrameFunctions = runtimeMirror(getClass.getClassLoader) + .reflect(functions) + .symbol + .typeSignature + .decls + .filter(s => s.isMethod && s.isPublic) + .map(_.name.toString) .toSet .filter(_.matches(word_pattern)) .diff(excludedDataFrameFunctions) From 398457af59875120ea8b3ed44468a51597e6a441 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Thu, 19 Sep 2024 09:02:34 -0400 Subject: [PATCH 073/189] [SPARK-49422][CONNECT][SQL] Add groupByKey to sql/api ### What changes were proposed in this pull request? This PR adds `Dataset.groupByKey(..)` to the shared interface. I forgot to add in the previous PR. ### Why are the changes needed? The shared interface needs to support all functionality. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48147 from hvanhovell/SPARK-49422-follow-up. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../scala/org/apache/spark/sql/Dataset.scala | 24 ++----- .../org/apache/spark/sql/api/Dataset.scala | 22 ++++++ .../scala/org/apache/spark/sql/Dataset.scala | 68 +++---------------- 3 files changed, 39 insertions(+), 75 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index 161a0d9d265f0..accfff9f2b073 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -524,27 +524,11 @@ class Dataset[T] private[sql] ( result(0) } - /** - * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 3.5.0 - */ + /** @inheritdoc */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { KeyValueGroupedDatasetImpl[K, T](this, encoderFor[K], func) } - /** - * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given - * key `func`. - * - * @group typedrel - * @since 3.5.0 - */ - def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - groupByKey(ToScalaUDF(func))(encoder) - /** @inheritdoc */ @scala.annotation.varargs def rollup(cols: Column*): RelationalGroupedDataset = { @@ -1480,4 +1464,10 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) + + /** @inheritdoc */ + override def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala index 284a69fe6ee3e..6eef034aa5157 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala @@ -1422,6 +1422,28 @@ abstract class Dataset[T] extends Serializable { */ def reduce(func: ReduceFunction[T]): T = reduce(ToScalaUDF(func)) + /** + * (Scala-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 2.0.0 + */ + def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] + + /** + * (Java-specific) Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given + * key `func`. + * + * @group typedrel + * @since 2.0.0 + */ + def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = { + groupByKey(ToScalaUDF(func))(encoder) + } + /** * Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns * set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, except for the aggregation, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 61f9e6ff7c042..ef628ca612b49 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -62,7 +62,7 @@ import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, FileTable} import org.apache.spark.sql.execution.python.EvaluatePython import org.apache.spark.sql.execution.stat.StatFunctions -import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf, ToScalaUDF} +import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf} import org.apache.spark.sql.internal.ExpressionUtils.column import org.apache.spark.sql.internal.TypedAggUtils.withInputType import org.apache.spark.sql.streaming.DataStreamWriter @@ -865,24 +865,7 @@ class Dataset[T] private[sql]( Filter(condition.expr, logicalPlan) } - /** - * Groups the Dataset using the specified columns, so we can run aggregation on them. See - * [[RelationalGroupedDataset]] for all the available aggregate functions. - * - * {{{ - * // Compute the average for all numeric columns grouped by department. - * ds.groupBy($"department").avg() - * - * // Compute the max age and average salary, grouped by department and gender. - * ds.groupBy($"department", $"gender").agg(Map( - * "salary" -> "avg", - * "age" -> "max" - * )) - * }}} - * - * @group untypedrel - * @since 2.0.0 - */ + /** @inheritdoc */ @scala.annotation.varargs def groupBy(cols: Column*): RelationalGroupedDataset = { RelationalGroupedDataset(toDF(), cols.map(_.expr), RelationalGroupedDataset.GroupByType) @@ -914,13 +897,7 @@ class Dataset[T] private[sql]( rdd.reduce(func) } - /** - * (Scala-specific) - * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. - * - * @group typedrel - * @since 2.0.0 - */ + /** @inheritdoc */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { val withGroupingKey = AppendColumns(func, logicalPlan) val executed = sparkSession.sessionState.executePlan(withGroupingKey) @@ -933,16 +910,6 @@ class Dataset[T] private[sql]( withGroupingKey.newColumns) } - /** - * (Java-specific) - * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`. - * - * @group typedrel - * @since 2.0.0 - */ - def groupByKey[K](func: MapFunction[T, K], encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = - groupByKey(ToScalaUDF(func))(encoder) - /** @inheritdoc */ def unpivot( ids: Array[Column], @@ -1640,28 +1607,7 @@ class Dataset[T] private[sql]( new DataFrameWriterV2Impl[T](table, this) } - /** - * Merges a set of updates, insertions, and deletions based on a source table into - * a target table. - * - * Scala Examples: - * {{{ - * spark.table("source") - * .mergeInto("target", $"source.id" === $"target.id") - * .whenMatched($"salary" === 100) - * .delete() - * .whenNotMatched() - * .insertAll() - * .whenNotMatchedBySource($"salary" === 100) - * .update(Map( - * "salary" -> lit(200) - * )) - * .merge() - * }}} - * - * @group basic - * @since 4.0.0 - */ + /** @inheritdoc */ def mergeInto(table: String, condition: Column): MergeIntoWriter[T] = { if (isStreaming) { logicalPlan.failAnalysis( @@ -2024,6 +1970,12 @@ class Dataset[T] private[sql]( @scala.annotation.varargs override def agg(expr: Column, exprs: Column*): DataFrame = super.agg(expr, exprs: _*) + /** @inheritdoc */ + override def groupByKey[K]( + func: MapFunction[T, K], + encoder: Encoder[K]): KeyValueGroupedDataset[K, T] = + super.groupByKey(func, encoder).asInstanceOf[KeyValueGroupedDataset[K, T]] + //////////////////////////////////////////////////////////////////////////// // For Python API //////////////////////////////////////////////////////////////////////////// From 94dca78c128ff3d1571326629b4100ee092afb54 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 21:10:52 +0800 Subject: [PATCH 074/189] [SPARK-49693][PYTHON][CONNECT] Refine the string representation of `timedelta` ### What changes were proposed in this pull request? Refine the string representation of `timedelta`, by following the ISO format. Note that the used units in JVM side (`Duration`) and Pandas are different. ### Why are the changes needed? We should not leak the raw data ### Does this PR introduce _any_ user-facing change? yes PySpark Classic: ``` In [1]: from pyspark.sql import functions as sf In [2]: import datetime In [3]: sf.lit(datetime.timedelta(1, 1)) Out[3]: Column<'PT24H1S'> ``` PySpark Connect (before): ``` In [1]: from pyspark.sql import functions as sf In [2]: import datetime In [3]: sf.lit(datetime.timedelta(1, 1)) Out[3]: Column<'86401000000'> ``` PySpark Connect (after): ``` In [1]: from pyspark.sql import functions as sf In [2]: import datetime In [3]: sf.lit(datetime.timedelta(1, 1)) Out[3]: Column<'P1DT0H0M1S'> ``` ### How was this patch tested? added test ### Was this patch authored or co-authored using generative AI tooling? no Closes #48159 from zhengruifeng/pc_lit_delta. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- python/pyspark/sql/connect/expressions.py | 12 +++++++++++- python/pyspark/sql/tests/test_column.py | 23 ++++++++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/connect/expressions.py b/python/pyspark/sql/connect/expressions.py index 63128ef48e389..0b5512b61925c 100644 --- a/python/pyspark/sql/connect/expressions.py +++ b/python/pyspark/sql/connect/expressions.py @@ -489,7 +489,17 @@ def __repr__(self) -> str: ts = TimestampNTZType().fromInternal(self._value) if ts is not None and isinstance(ts, datetime.datetime): return ts.strftime("%Y-%m-%d %H:%M:%S.%f") - # TODO(SPARK-49693): Refine the string representation of timedelta + elif isinstance(self._dataType, DayTimeIntervalType): + delta = DayTimeIntervalType().fromInternal(self._value) + if delta is not None and isinstance(delta, datetime.timedelta): + import pandas as pd + + # Note: timedelta itself does not provide isoformat method. + # Both Pandas and java.time.Duration provide it, but the format + # is sightly different: + # java.time.Duration only applies HOURS, MINUTES, SECONDS units, + # while Pandas applies all supported units. + return pd.Timedelta(delta).isoformat() # type: ignore[attr-defined] return f"{self._value}" diff --git a/python/pyspark/sql/tests/test_column.py b/python/pyspark/sql/tests/test_column.py index 220ecd387f7ee..1972dd2804d98 100644 --- a/python/pyspark/sql/tests/test_column.py +++ b/python/pyspark/sql/tests/test_column.py @@ -19,12 +19,13 @@ from enum import Enum from itertools import chain import datetime +import unittest from pyspark.sql import Column, Row from pyspark.sql import functions as sf from pyspark.sql.types import StructType, StructField, IntegerType, LongType from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError -from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, pandas_requirement_message class ColumnTestsMixin: @@ -289,6 +290,26 @@ def test_lit_time_representation(self): ts = datetime.datetime(2021, 3, 4, 12, 34, 56, 1234) self.assertEqual(str(sf.lit(ts)), "Column<'2021-03-04 12:34:56.001234'>") + @unittest.skipIf(not have_pandas, pandas_requirement_message) + def test_lit_delta_representation(self): + for delta in [ + datetime.timedelta(days=1), + datetime.timedelta(hours=2), + datetime.timedelta(minutes=3), + datetime.timedelta(seconds=4), + datetime.timedelta(microseconds=5), + datetime.timedelta(days=2, hours=21, microseconds=908), + datetime.timedelta(days=1, minutes=-3, microseconds=-1001), + datetime.timedelta(days=1, hours=2, minutes=3, seconds=4, microseconds=5), + ]: + import pandas as pd + + # Column<'PT69H0.000908S'> or Column<'P2DT21H0M0.000908S'> + s = str(sf.lit(delta)) + + # Parse the ISO string representation and compare + self.assertTrue(pd.Timedelta(s[8:-2]).to_pytimedelta() == delta) + def test_enum_literals(self): class IntEnum(Enum): X = 1 From f0fb0c89ec29b587569d68a824c4ce7543721c06 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 10:06:45 -0700 Subject: [PATCH 075/189] [SPARK-49719][SQL] Make `UUID` and `SHUFFLE` accept integer `seed` ### What changes were proposed in this pull request? Make `UUID` and `SHUFFLE` accept integer `seed` ### Why are the changes needed? In most cases, `seed` accept both int and long, but `UUID` and `SHUFFLE` only accept long seed ```py In [1]: spark.sql("SELECT RAND(1L), RAND(1), SHUFFLE(array(1, 20, 3, 5), 1L), UUID(1L)").show() +------------------+------------------+---------------------------+--------------------+ | rand(1)| rand(1)|shuffle(array(1, 20, 3, 5))| uuid()| +------------------+------------------+---------------------------+--------------------+ |0.6363787615254752|0.6363787615254752| [20, 1, 3, 5]|1ced31d7-59ef-4bb...| +------------------+------------------+---------------------------+--------------------+ In [2]: spark.sql("SELECT UUID(1)").show() ... AnalysisException: [INVALID_PARAMETER_VALUE.LONG] The value of parameter(s) `seed` in `UUID` is invalid: expects a long literal, but got "1". SQLSTATE: 22023; line 1 pos 7 ... In [3]: spark.sql("SELECT SHUFFLE(array(1, 20, 3, 5), 1)").show() ... AnalysisException: [INVALID_PARAMETER_VALUE.LONG] The value of parameter(s) `seed` in `shuffle` is invalid: expects a long literal, but got "1". SQLSTATE: 22023; line 1 pos 7 ... ``` ### Does this PR introduce _any_ user-facing change? yes after this fix: ```py In [2]: spark.sql("SELECT SHUFFLE(array(1, 20, 3, 5), 1L), SHUFFLE(array(1, 20, 3, 5), 1), UUID(1L), UUID(1)").show() +---------------------------+---------------------------+--------------------+--------------------+ |shuffle(array(1, 20, 3, 5))|shuffle(array(1, 20, 3, 5))| uuid()| uuid()| +---------------------------+---------------------------+--------------------+--------------------+ | [20, 1, 3, 5]| [20, 1, 3, 5]|1ced31d7-59ef-4bb...|1ced31d7-59ef-4bb...| +---------------------------+---------------------------+--------------------+--------------------+ ``` ### How was this patch tested? added tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48166 from zhengruifeng/int_seed. Authored-by: Ruifeng Zheng Signed-off-by: Dongjoon Hyun --- .../sql/catalyst/expressions/randomExpressions.scala | 1 + .../catalyst/expressions/CollectionExpressionsSuite.scala | 8 ++++++++ .../sql/catalyst/expressions/MiscExpressionsSuite.scala | 7 +++++++ 3 files changed, 16 insertions(+) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala index ea9ca451c2cb1..f329f8346b0de 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala @@ -81,6 +81,7 @@ trait ExpressionWithRandomSeed extends Expression { private[catalyst] object ExpressionWithRandomSeed { def expressionToSeed(e: Expression, source: String): Option[Long] = e match { + case IntegerLiteral(seed) => Some(seed) case LongLiteral(seed) => Some(seed) case Literal(null, _) => None case _ => throw QueryCompilationErrors.invalidRandomSeedParameter(source, e) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala index e9de59b3ec48c..55148978fa005 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala @@ -2293,6 +2293,14 @@ class CollectionExpressionsSuite evaluateWithMutableProjection(Shuffle(ai0, seed2))) assert(evaluateWithUnsafeProjection(Shuffle(ai0, seed1)) !== evaluateWithUnsafeProjection(Shuffle(ai0, seed2))) + + val seed3 = Literal.create(r.nextInt()) + assert(evaluateWithoutCodegen(new Shuffle(ai0, seed3)) === + evaluateWithoutCodegen(new Shuffle(ai0, seed3))) + assert(evaluateWithMutableProjection(new Shuffle(ai0, seed3)) === + evaluateWithMutableProjection(new Shuffle(ai0, seed3))) + assert(evaluateWithUnsafeProjection(new Shuffle(ai0, seed3)) === + evaluateWithUnsafeProjection(new Shuffle(ai0, seed3))) } test("Array Except") { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala index 1f37886f44258..40e6fe1a90a63 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala @@ -71,6 +71,13 @@ class MiscExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { evaluateWithMutableProjection(Uuid(seed2))) assert(evaluateWithUnsafeProjection(Uuid(seed1)) !== evaluateWithUnsafeProjection(Uuid(seed2))) + + val seed3 = Literal.create(r.nextInt()) + assert(evaluateWithoutCodegen(new Uuid(seed3)) === evaluateWithoutCodegen(new Uuid(seed3))) + assert(evaluateWithMutableProjection(new Uuid(seed3)) === + evaluateWithMutableProjection(new Uuid(seed3))) + assert(evaluateWithUnsafeProjection(new Uuid(seed3)) === + evaluateWithUnsafeProjection(new Uuid(seed3))) } test("PrintToStderr") { From 92cad2abd54e775259dc36d2f90242460d72a174 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Thu, 19 Sep 2024 10:09:36 -0700 Subject: [PATCH 076/189] [SPARK-49716][PS][DOCS][TESTS] Fix documentation and add test of barh plot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? - Update the documentation for barh plot to clarify the difference between axis interpretation in Plotly and Matplotlib. - Test multiple columns as value axis. The parameter difference is demonstrated as below. ```py >>> df = ps.DataFrame({'lab': ['A', 'B', 'C'], 'val': [10, 30, 20]}) >>> df.plot.barh(x='val', y='lab').show() # plot1 >>> ps.set_option('plotting.backend', 'matplotlib') >>> import matplotlib.pyplot as plt >>> df.plot.barh(x='lab', y='val') >>> plt.show() # plot2 ``` plot1 ![newplot (5)](https://github.com/user-attachments/assets/f1b6fabe-9509-41bb-8cfb-0733f65f1643) plot2 ![Figure_1](https://github.com/user-attachments/assets/10e1b65f-6116-4490-9956-29e1fbf0c053) ### Why are the changes needed? The barh plot’s x and y axis behavior differs between Plotly and Matplotlib, which may confuse users. The updated documentation and tests help ensure clarity and prevent misinterpretation. ### Does this PR introduce _any_ user-facing change? No. Doc change only. ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48161 from xinrong-meng/ps_barh. Authored-by: Xinrong Meng Signed-off-by: Dongjoon Hyun --- python/pyspark/pandas/plot/core.py | 13 ++++++++++--- .../pandas/tests/plot/test_frame_plot_plotly.py | 5 +++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index 7630ecc398954..429e97ecf07bb 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -756,10 +756,10 @@ def barh(self, x=None, y=None, **kwargs): Parameters ---------- - x : label or position, default DataFrame.index - Column to be used for categories. - y : label or position, default All numeric columns in dataframe + x : label or position, default All numeric columns in dataframe Columns to be plotted from the DataFrame. + y : label or position, default DataFrame.index + Column to be used for categories. **kwds Keyword arguments to pass on to :meth:`pyspark.pandas.DataFrame.plot` or :meth:`pyspark.pandas.Series.plot`. @@ -770,6 +770,13 @@ def barh(self, x=None, y=None, **kwargs): Return an custom object when ``backend!=plotly``. Return an ndarray when ``subplots=True`` (matplotlib-only). + Notes + ----- + In Plotly and Matplotlib, the interpretation of `x` and `y` for `barh` plots differs. + In Plotly, `x` refers to the values and `y` refers to the categories. + In Matplotlib, `x` refers to the categories and `y` refers to the values. + Ensure correct axis labeling based on the backend used. + See Also -------- plotly.express.bar : Plot a vertical bar plot using plotly. diff --git a/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py b/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py index 37469db2c8f51..8d197649aaebe 100644 --- a/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py @@ -105,9 +105,10 @@ def check_barh_plot_with_x_y(pdf, psdf, x, y): self.assertEqual(pdf.plot.barh(x=x, y=y), psdf.plot.barh(x=x, y=y)) # this is testing plot with specified x and y - pdf1 = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20]}) + pdf1 = pd.DataFrame({"lab": ["A", "B", "C"], "val": [10, 30, 20], "val2": [1.1, 2.2, 3.3]}) psdf1 = ps.from_pandas(pdf1) - check_barh_plot_with_x_y(pdf1, psdf1, x="lab", y="val") + check_barh_plot_with_x_y(pdf1, psdf1, x="val", y="lab") + check_barh_plot_with_x_y(pdf1, psdf1, x=["val", "val2"], y="lab") def test_barh_plot(self): def check_barh_plot(pdf, psdf): From 6d1815eceea2003de2e3602f0f64e8188e8288d8 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 12:31:48 -0700 Subject: [PATCH 077/189] [SPARK-49718][PS] Switch `Scatter` plot to sampled data ### What changes were proposed in this pull request? Switch `Scatter` plot to sampled data ### Why are the changes needed? when the data distribution has relationship with the order, the first n rows will not be representative of the whole dataset for example: ``` import pandas as pd import numpy as np import pyspark.pandas as ps # ps.set_option("plotting.max_rows", 10000) np.random.seed(123) pdf = pd.DataFrame(np.random.randn(10000, 4), columns=list('ABCD')).sort_values("A") psdf = ps.DataFrame(pdf) psdf.plot.scatter(x='B', y='A') ``` all 10k datapoints: ![image](https://github.com/user-attachments/assets/72cf7e97-ad10-41e0-a8a6-351747d5285f) before (first 1k datapoints): ![image](https://github.com/user-attachments/assets/1ed50d2c-7772-4579-a84c-6062542d9367) after (sampled 1k datapoints): ![image](https://github.com/user-attachments/assets/6c684cba-4119-4c38-8228-2bedcdeb9e59) ### Does this PR introduce _any_ user-facing change? yes ### How was this patch tested? ci and manually test ### Was this patch authored or co-authored using generative AI tooling? no Closes #48164 from zhengruifeng/ps_scatter_sampling. Authored-by: Ruifeng Zheng Signed-off-by: Dongjoon Hyun --- python/pyspark/pandas/plot/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index 429e97ecf07bb..6f036b7669246 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -479,7 +479,7 @@ class PandasOnSparkPlotAccessor(PandasObject): "pie": TopNPlotBase().get_top_n, "bar": TopNPlotBase().get_top_n, "barh": TopNPlotBase().get_top_n, - "scatter": TopNPlotBase().get_top_n, + "scatter": SampledPlotBase().get_sampled, "area": SampledPlotBase().get_sampled, "line": SampledPlotBase().get_sampled, } From 04455797bfb3631b13b41cfa5d2604db3bf8acc2 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 19 Sep 2024 12:32:30 -0700 Subject: [PATCH 078/189] [SPARK-49720][PYTHON][INFRA] Add a script to clean up PySpark temp files ### What changes were proposed in this pull request? Add a script to clean up PySpark temp files ### Why are the changes needed? Sometimes I encounter weird issues due to the out-dated `pyspark.zip` file, and removing it can result in expected behavior. So I think we can add such a script. ### Does this PR introduce _any_ user-facing change? no, dev-only ### How was this patch tested? manually test ### Was this patch authored or co-authored using generative AI tooling? no Closes #48167 from zhengruifeng/py_infra_cleanup. Authored-by: Ruifeng Zheng Signed-off-by: Dongjoon Hyun --- dev/py-cleanup | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100755 dev/py-cleanup diff --git a/dev/py-cleanup b/dev/py-cleanup new file mode 100755 index 0000000000000..6a2edd1040171 --- /dev/null +++ b/dev/py-cleanup @@ -0,0 +1,31 @@ +#!/usr/bin/env bash + +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Utility for temporary files cleanup in 'python'. +# usage: ./dev/py-cleanup + +set -ex + +SPARK_HOME="$(cd "`dirname $0`"/..; pwd)" +cd "$SPARK_HOME" + +rm -rf python/target +rm -rf python/lib/pyspark.zip +rm -rf python/docs/build +rm -rf python/docs/source/reference/*/api From ca726c10925a3677bf057f65ecf415e608c63cd5 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Thu, 19 Sep 2024 17:16:25 -0700 Subject: [PATCH 079/189] [SPARK-49721][BUILD] Upgrade `protobuf-java` to 3.25.5 ### What changes were proposed in this pull request? This PR aims to upgrade `protobuf-java` to 3.25.5. ### Why are the changes needed? To bring the latest bug fixes. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48170 Closes #48171 from dongjoon-hyun/SPARK-49721. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- pom.xml | 2 +- project/SparkBuild.scala | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pom.xml b/pom.xml index 694ea31e6f377..ddabc82d2ad13 100644 --- a/pom.xml +++ b/pom.xml @@ -124,7 +124,7 @@ 3.4.0 - 3.25.4 + 3.25.5 3.11.4 ${hadoop.version} 3.9.2 diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index d93a52985b772..2f390cb70baa8 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -89,7 +89,7 @@ object BuildCommons { // Google Protobuf version used for generating the protobuf. // SPARK-41247: needs to be consistent with `protobuf.version` in `pom.xml`. - val protoVersion = "3.25.4" + val protoVersion = "3.25.5" // GRPC version used for Spark Connect. val grpcVersion = "1.62.2" } From a5ac80af8e94afe56105c265a94d02ef878e1de9 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 20 Sep 2024 08:29:48 +0800 Subject: [PATCH 080/189] [SPARK-49713][PYTHON][CONNECT] Make function `count_min_sketch` accept number arguments ### What changes were proposed in this pull request? 1, Make function `count_min_sketch` accept number arguments; 2, Make argument `seed` optional; 3, fix the type hints of `eps/confidence/seed` from `ColumnOrName` to `Column`, because they require a foldable value and actually do not accept column name: ``` In [3]: from pyspark.sql import functions as sf In [4]: df = spark.range(10000).withColumn("seed", sf.lit(1).cast("int")) In [5]: df.select(sf.hex(sf.count_min_sketch("id", sf.lit(0.5), sf.lit(0.5), "seed"))) ... AnalysisException: [DATATYPE_MISMATCH.NON_FOLDABLE_INPUT] Cannot resolve "count_min_sketch(id, 0.5, 0.5, seed)" due to data type mismatch: the input `seed` should be a foldable "INT" expression; however, got "seed". SQLSTATE: 42K09; 'Aggregate [unresolvedalias('hex(count_min_sketch(id#1L, 0.5, 0.5, seed#2, 0, 0)))] +- Project [id#1L, cast(1 as int) AS seed#2] +- Range (0, 10000, step=1, splits=Some(12)) ... ``` ### Why are the changes needed? 1, seed is optional in other similar functions; 2, existing type hint is `ColumnOrName` which is misleading since column name is not actually supported ### Does this PR introduce _any_ user-facing change? yes, it support number arguments ### How was this patch tested? updated doctests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48157 from zhengruifeng/py_fix_count_min_sketch. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- .../pyspark/sql/connect/functions/builtin.py | 10 +-- python/pyspark/sql/functions/builtin.py | 71 +++++++++++++++---- .../org/apache/spark/sql/functions.scala | 12 ++++ 3 files changed, 77 insertions(+), 16 deletions(-) diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index 2870d9c408b6b..7fed175cbc8ea 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -71,6 +71,7 @@ StringType, ) from pyspark.sql.utils import enum_to_value as _enum_to_value +from pyspark.util import JVM_INT_MAX # The implementation of pandas_udf is embedded in pyspark.sql.function.pandas_udf # for code reuse. @@ -1126,11 +1127,12 @@ def grouping_id(*cols: "ColumnOrName") -> Column: def count_min_sketch( col: "ColumnOrName", - eps: "ColumnOrName", - confidence: "ColumnOrName", - seed: "ColumnOrName", + eps: Union[Column, float], + confidence: Union[Column, float], + seed: Optional[Union[Column, int]] = None, ) -> Column: - return _invoke_function_over_columns("count_min_sketch", col, eps, confidence, seed) + _seed = lit(random.randint(0, JVM_INT_MAX)) if seed is None else lit(seed) + return _invoke_function_over_columns("count_min_sketch", col, lit(eps), lit(confidence), _seed) count_min_sketch.__doc__ = pysparkfuncs.count_min_sketch.__doc__ diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index c0730b193bc72..5f8d1c21a24f1 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -6015,9 +6015,9 @@ def grouping_id(*cols: "ColumnOrName") -> Column: @_try_remote_functions def count_min_sketch( col: "ColumnOrName", - eps: "ColumnOrName", - confidence: "ColumnOrName", - seed: "ColumnOrName", + eps: Union[Column, float], + confidence: Union[Column, float], + seed: Optional[Union[Column, int]] = None, ) -> Column: """ Returns a count-min sketch of a column with the given esp, confidence and seed. @@ -6031,13 +6031,24 @@ def count_min_sketch( ---------- col : :class:`~pyspark.sql.Column` or str target column to compute on. - eps : :class:`~pyspark.sql.Column` or str + eps : :class:`~pyspark.sql.Column` or float relative error, must be positive - confidence : :class:`~pyspark.sql.Column` or str + + .. versionchanged:: 4.0.0 + `eps` now accepts float value. + + confidence : :class:`~pyspark.sql.Column` or float confidence, must be positive and less than 1.0 - seed : :class:`~pyspark.sql.Column` or str + + .. versionchanged:: 4.0.0 + `confidence` now accepts float value. + + seed : :class:`~pyspark.sql.Column` or int, optional random seed + .. versionchanged:: 4.0.0 + `seed` now accepts int value. + Returns ------- :class:`~pyspark.sql.Column` @@ -6045,12 +6056,48 @@ def count_min_sketch( Examples -------- - >>> df = spark.createDataFrame([[1], [2], [1]], ['data']) - >>> df = df.agg(count_min_sketch(df.data, lit(0.5), lit(0.5), lit(1)).alias('sketch')) - >>> df.select(hex(df.sketch).alias('r')).collect() - [Row(r='0000000100000000000000030000000100000004000000005D8D6AB90000000000000000000000000000000200000000000000010000000000000000')] - """ - return _invoke_function_over_columns("count_min_sketch", col, eps, confidence, seed) + Example 1: Using columns as arguments + + >>> from pyspark.sql import functions as sf + >>> spark.range(100).select( + ... sf.hex(sf.count_min_sketch(sf.col("id"), sf.lit(3.0), sf.lit(0.1), sf.lit(1))) + ... ).show(truncate=False) + +------------------------------------------------------------------------+ + |hex(count_min_sketch(id, 3.0, 0.1, 1)) | + +------------------------------------------------------------------------+ + |0000000100000000000000640000000100000001000000005D8D6AB90000000000000064| + +------------------------------------------------------------------------+ + + Example 2: Using numbers as arguments + + >>> from pyspark.sql import functions as sf + >>> spark.range(100).select( + ... sf.hex(sf.count_min_sketch("id", 1.0, 0.3, 2)) + ... ).show(truncate=False) + +----------------------------------------------------------------------------------------+ + |hex(count_min_sketch(id, 1.0, 0.3, 2)) | + +----------------------------------------------------------------------------------------+ + |0000000100000000000000640000000100000002000000005D96391C00000000000000320000000000000032| + +----------------------------------------------------------------------------------------+ + + Example 3: Using a random seed + + >>> from pyspark.sql import functions as sf + >>> spark.range(100).select( + ... sf.hex(sf.count_min_sketch("id", sf.lit(1.5), 0.6)) + ... ).show(truncate=False) # doctest: +SKIP + +----------------------------------------------------------------------------------------------------------------------------------------+ + |hex(count_min_sketch(id, 1.5, 0.6, 2120704260)) | + +----------------------------------------------------------------------------------------------------------------------------------------+ + |0000000100000000000000640000000200000002000000005ADECCEE00000000153EBE090000000000000033000000000000003100000000000000320000000000000032| + +----------------------------------------------------------------------------------------------------------------------------------------+ + """ # noqa: E501 + _eps = lit(eps) + _conf = lit(confidence) + if seed is None: + return _invoke_function_over_columns("count_min_sketch", col, _eps, _conf) + else: + return _invoke_function_over_columns("count_min_sketch", col, _eps, _conf, lit(seed)) @_try_remote_functions diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 02669270c8acf..0662b8f2b271f 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -389,6 +389,18 @@ object functions { def count_min_sketch(e: Column, eps: Column, confidence: Column, seed: Column): Column = Column.fn("count_min_sketch", e, eps, confidence, seed) + /** + * Returns a count-min sketch of a column with the given esp, confidence and seed. The result is + * an array of bytes, which can be deserialized to a `CountMinSketch` before usage. Count-min + * sketch is a probabilistic data structure used for cardinality estimation using sub-linear + * space. + * + * @group agg_funcs + * @since 4.0.0 + */ + def count_min_sketch(e: Column, eps: Column, confidence: Column): Column = + count_min_sketch(e, eps, confidence, lit(SparkClassUtils.random.nextInt)) + private[spark] def collect_top_k(e: Column, num: Int, reverse: Boolean): Column = Column.internalFn("collect_top_k", e, lit(num), lit(reverse)) From d4665fa1df716305acb49912d41c396b39343c93 Mon Sep 17 00:00:00 2001 From: Anish Shrigondekar Date: Fri, 20 Sep 2024 14:11:14 +0900 Subject: [PATCH 081/189] [SPARK-49677][SS] Ensure that changelog files are written on commit and forceSnapshot flag is also reset ### What changes were proposed in this pull request? Ensure that changelog files are written on commit and forceSnapshot flag is also reset ### Why are the changes needed? Without these changes, we are not writing the changelog files per batch and we are also trying to upload full snapshot each time since the flag is not being reset correctly ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Added unit tests Before: ``` [info] Run completed in 3 seconds, 438 milliseconds. [info] Total number of tests run: 1 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 0, failed 1, canceled 0, ignored 0, pending 0 [info] *** 1 TEST FAILED *** ``` After: ``` [info] Run completed in 4 seconds, 155 milliseconds. [info] Total number of tests run: 1 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 1, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48125 from anishshri-db/task/SPARK-49677. Authored-by: Anish Shrigondekar Signed-off-by: Jungtaek Lim --- .../execution/streaming/state/RocksDB.scala | 16 ++++---- .../streaming/state/RocksDBSuite.scala | 41 +++++++++++++++++++ 2 files changed, 49 insertions(+), 8 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala index 81e80629092a0..4a2aac43b3331 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala @@ -646,15 +646,15 @@ class RocksDB( // is enabled. if (shouldForceSnapshot.get()) { uploadSnapshot() + shouldForceSnapshot.set(false) + } + + // ensure that changelog files are always written + try { + assert(changelogWriter.isDefined) + changelogWriter.foreach(_.commit()) + } finally { changelogWriter = None - changelogWriter.foreach(_.abort()) - } else { - try { - assert(changelogWriter.isDefined) - changelogWriter.foreach(_.commit()) - } finally { - changelogWriter = None - } } } else { assert(changelogWriter.isEmpty) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala index 691f18451af22..608a22a284b6c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala @@ -811,6 +811,47 @@ class RocksDBSuite extends AlsoTestWithChangelogCheckpointingEnabled with Shared } } + testWithChangelogCheckpointingEnabled("RocksDB: ensure that changelog files are written " + + "and snapshots uploaded optionally with changelog format v2") { + withTempDir { dir => + val remoteDir = Utils.createTempDir().toString + val conf = dbConf.copy(minDeltasForSnapshot = 5, compactOnCommit = false) + new File(remoteDir).delete() // to make sure that the directory gets created + withDB(remoteDir, conf = conf, useColumnFamilies = true) { db => + db.createColFamilyIfAbsent("test") + db.load(0) + db.put("a", "1") + db.put("b", "2") + db.commit() + assert(changelogVersionsPresent(remoteDir) == Seq(1)) + assert(snapshotVersionsPresent(remoteDir) == Seq(1)) + + db.load(1) + db.put("a", "3") + db.put("c", "4") + db.commit() + + assert(changelogVersionsPresent(remoteDir) == Seq(1, 2)) + assert(snapshotVersionsPresent(remoteDir) == Seq(1)) + + db.removeColFamilyIfExists("test") + db.load(2) + db.remove("a") + db.put("d", "5") + db.commit() + assert(changelogVersionsPresent(remoteDir) == Seq(1, 2, 3)) + assert(snapshotVersionsPresent(remoteDir) == Seq(1, 3)) + + db.load(3) + db.put("e", "6") + db.remove("b") + db.commit() + assert(changelogVersionsPresent(remoteDir) == Seq(1, 2, 3, 4)) + assert(snapshotVersionsPresent(remoteDir) == Seq(1, 3)) + } + } + } + test("RocksDB: ensure merge operation correctness") { withTempDir { dir => val remoteDir = Utils.createTempDir().toString From 6352c12f607bc092c33f1f29174d6699f8312380 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Fri, 20 Sep 2024 15:29:08 +0900 Subject: [PATCH 082/189] [MINOR][INFRA] Disable 'pages build and deployment' action ### What changes were proposed in this pull request? Disable https://github.com/apache/spark/actions/runs/10951008649/ via: > adding a .nojekyll file to the root of your source branch will bypass the Jekyll build process and deploy the content directly. https://docs.github.com/en/pages/quickstart ### Why are the changes needed? restore ci ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? no ### Was this patch authored or co-authored using generative AI tooling? no Closes #48176 from yaooqinn/action. Authored-by: Kent Yao Signed-off-by: Hyukjin Kwon --- .nojekyll | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 .nojekyll diff --git a/.nojekyll b/.nojekyll new file mode 100644 index 0000000000000..e69de29bb2d1d From c009cd061c4923955a1e7ec9bf6c045f93d27ef7 Mon Sep 17 00:00:00 2001 From: Uros Bojanic Date: Fri, 20 Sep 2024 09:16:04 +0200 Subject: [PATCH 083/189] [SPARK-49392][SQL][FOLLOWUP] Catch errors when failing to write to external data source ### What changes were proposed in this pull request? Change `sqlState` to KD010. ### Why are the changes needed? Necessary modification for the Databricks error class space. ### Does this PR introduce _any_ user-facing change? Yes, the new error message is now updated to KD010. ### How was this patch tested? Existing tests (updated). ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48165 from uros-db/external-data-source-fix. Authored-by: Uros Bojanic Signed-off-by: Max Gekk --- common/utils/src/main/resources/error/error-conditions.json | 2 +- common/utils/src/main/resources/error/error-states.json | 2 +- .../apache/spark/sql/errors/QueryCompilationErrorsSuite.scala | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 72985de6631f0..e83202d9e5ee3 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -1055,7 +1055,7 @@ "message" : [ "Encountered error when saving to external data source." ], - "sqlState" : "KD00F" + "sqlState" : "KD010" }, "DATA_SOURCE_NOT_EXIST" : { "message" : [ diff --git a/common/utils/src/main/resources/error/error-states.json b/common/utils/src/main/resources/error/error-states.json index edba6e1d43216..87811fef9836e 100644 --- a/common/utils/src/main/resources/error/error-states.json +++ b/common/utils/src/main/resources/error/error-states.json @@ -7417,7 +7417,7 @@ "standard": "N", "usedBy": ["Databricks"] }, - "KD00F": { + "KD010": { "description": "external data source failure", "origin": "Databricks", "standard": "N", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala index 370c118de9a93..832e1873af6a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala @@ -941,7 +941,7 @@ class QueryCompilationErrorsSuite cmd.run(spark) }, condition = "DATA_SOURCE_EXTERNAL_ERROR", - sqlState = "KD00F", + sqlState = "KD010", parameters = Map.empty ) } From b37863d2327131c670fe791576a907bcb5243cd6 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Fri, 20 Sep 2024 16:40:36 +0900 Subject: [PATCH 084/189] [MINOR][FOLLOWUP] Fix rat check for .nojekyll ### What changes were proposed in this pull request? Fix rat check for .nojekyll ### Why are the changes needed? CI fix ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? dev/check-license Ignored 1 lines in your exclusion files as comments or empty lines. RAT checks passed. ### Was this patch authored or co-authored using generative AI tooling? no Closes #48178 from yaooqinn/f. Authored-by: Kent Yao Signed-off-by: Hyukjin Kwon --- dev/.rat-excludes | 1 + 1 file changed, 1 insertion(+) diff --git a/dev/.rat-excludes b/dev/.rat-excludes index f38fd7e2012a5..b82cb7078c9f3 100644 --- a/dev/.rat-excludes +++ b/dev/.rat-excludes @@ -140,3 +140,4 @@ ui-test/package.json ui-test/package-lock.json core/src/main/resources/org/apache/spark/ui/static/package.json .*\.har +.nojekyll From 46b0210edb4ef8490ee4bbc4a40baf202a531b33 Mon Sep 17 00:00:00 2001 From: Nick Young Date: Fri, 20 Sep 2024 18:05:28 +0900 Subject: [PATCH 085/189] [SPARK-49699][SS] Disable PruneFilters for streaming workloads ### What changes were proposed in this pull request? The PR proposes to disable PruneFilters if the predicate of the filter is evaluated to `null` / `false` and the filter (and subtree) is streaming. ### Why are the changes needed? PruneFilters replaces the `null` / `false` filter with an empty relation, which means the subtree of the filter is also lost. The optimization does not care about whichever operator is in the subtree, hence some important operators like stateful operator, watermark node, observe node could be lost. The filter could be evaluated to `null` / `false` selectively among microbatches in various reasons (one simple example is the modification of the query during restart), which means stateful operator might not be available for batch N and be available for batch N + 1. For this case, streaming query will fail as batch N + 1 cannot load the state from batch N, and it's not recoverable in most cases. See new tests in StreamingQueryOptimizationCorrectnessSuite for details. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? UT. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48149 from n-young-db/n-young-db/disable-streaming-prune-filters. Lead-authored-by: Nick Young Co-authored-by: Jungtaek Lim Signed-off-by: Jungtaek Lim --- .../sql/catalyst/optimizer/Optimizer.scala | 7 +- .../apache/spark/sql/internal/SQLConf.scala | 9 +++ .../PropagateEmptyRelationSuite.scala | 27 ++++++-- .../optimizer/PruneFiltersSuite.scala | 34 ++++++++++ .../sql/execution/streaming/OffsetSeq.scala | 7 +- ...ingQueryOptimizationCorrectnessSuite.scala | 64 ++++++++++++++++++- 6 files changed, 137 insertions(+), 11 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 6ceeeb9bfdf38..8e14537c6a5b4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -1723,15 +1723,18 @@ object EliminateSorts extends Rule[LogicalPlan] { * 3) by eliminating the always-true conditions given the constraints on the child's output. */ object PruneFilters extends Rule[LogicalPlan] with PredicateHelper { + private def shouldApply(child: LogicalPlan): Boolean = + SQLConf.get.getConf(SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN) || !child.isStreaming + def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( _.containsPattern(FILTER), ruleId) { // If the filter condition always evaluate to true, remove the filter. case Filter(Literal(true, BooleanType), child) => child // If the filter condition always evaluate to null or false, // replace the input with an empty relation. - case Filter(Literal(null, _), child) => + case Filter(Literal(null, _), child) if shouldApply(child) => LocalRelation(child.output, data = Seq.empty, isStreaming = child.isStreaming) - case Filter(Literal(false, BooleanType), child) => + case Filter(Literal(false, BooleanType), child) if shouldApply(child) => LocalRelation(child.output, data = Seq.empty, isStreaming = child.isStreaming) // If any deterministic condition is guaranteed to be true given the constraints on the child's // output, remove the condition diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 094fb8f050bc8..2eaafde52228b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3827,6 +3827,15 @@ object SQLConf { .intConf .createWithDefault(ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) + val PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN = + buildConf("spark.databricks.sql.optimizer.pruneFiltersCanPruneStreamingSubplan") + .internal() + .doc("Allow PruneFilters to remove streaming subplans when we encounter a false filter. " + + "This flag is to restore prior buggy behavior for broken pipelines.") + .version("4.0.0") + .booleanConf + .createWithDefault(false) + object Deprecated { val MAPRED_REDUCE_TASKS = "mapred.reduce.tasks" } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelationSuite.scala index 5aeb27f7ee6b4..451236162343b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelationSuite.scala @@ -27,12 +27,13 @@ import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical.{Expand, Filter, LocalRelation, LogicalPlan, Project} import org.apache.spark.sql.catalyst.rules.RuleExecutor import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{IntegerType, MetadataBuilder} -class PropagateEmptyRelationSuite extends PlanTest { +class PropagateEmptyRelationSuite extends PlanTest { object Optimize extends RuleExecutor[LogicalPlan] { val batches = - Batch("PropagateEmptyRelation", Once, + Batch("PropagateEmptyRelation", FixedPoint(1), CombineUnions, ReplaceDistinctWithAggregate, ReplaceExceptWithAntiJoin, @@ -45,7 +46,7 @@ class PropagateEmptyRelationSuite extends PlanTest { object OptimizeWithoutPropagateEmptyRelation extends RuleExecutor[LogicalPlan] { val batches = - Batch("OptimizeWithoutPropagateEmptyRelation", Once, + Batch("OptimizeWithoutPropagateEmptyRelation", FixedPoint(1), CombineUnions, ReplaceDistinctWithAggregate, ReplaceExceptWithAntiJoin, @@ -216,10 +217,24 @@ class PropagateEmptyRelationSuite extends PlanTest { .where($"a" =!= 200) .orderBy($"a".asc) - val optimized = Optimize.execute(query.analyze) - val correctAnswer = LocalRelation(output, isStreaming = true) + withSQLConf( + SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> "true") { + val optimized = Optimize.execute(query.analyze) + val correctAnswer = LocalRelation(output, isStreaming = true) + comparePlans(optimized, correctAnswer) + } - comparePlans(optimized, correctAnswer) + withSQLConf( + SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> "false") { + val optimized = Optimize.execute(query.analyze) + val correctAnswer = relation + .where(false) + .where($"a" > 1) + .select($"a") + .where($"a" =!= 200) + .orderBy($"a".asc).analyze + comparePlans(optimized, correctAnswer) + } } test("SPARK-47305 correctly tag isStreaming when propagating empty relation " + diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PruneFiltersSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PruneFiltersSuite.scala index b81a57f4f8cd5..66ded338340f3 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PruneFiltersSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PruneFiltersSuite.scala @@ -174,4 +174,38 @@ class PruneFiltersSuite extends PlanTest { testRelation.where(!$"a".attr.in(1, 3, 5) && $"a".attr === 7 && $"b".attr === 1) .where(Rand(10) > 0.1 && Rand(10) < 1.1).analyze) } + + test("Streaming relation is not lost under true filter") { + Seq("true", "false").foreach(x => withSQLConf( + SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> x) { + val streamingRelation = + LocalRelation(Seq($"a".int, $"b".int, $"c".int), Nil, isStreaming = true) + val originalQuery = streamingRelation.where(10 > 5).select($"a").analyze + val optimized = Optimize.execute(originalQuery) + val correctAnswer = streamingRelation.select($"a").analyze + comparePlans(optimized, correctAnswer) + }) + } + + test("Streaming relation is not lost under false filter") { + withSQLConf( + SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> "true") { + val streamingRelation = + LocalRelation(Seq($"a".int, $"b".int, $"c".int), Nil, isStreaming = true) + val originalQuery = streamingRelation.where(10 < 5).select($"a").analyze + val optimized = Optimize.execute(originalQuery) + val correctAnswer = streamingRelation.select($"a").analyze + comparePlans(optimized, correctAnswer) + } + + withSQLConf( + SQLConf.PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> "false") { + val streamingRelation = + LocalRelation(Seq($"a".int, $"b".int, $"c".int), Nil, isStreaming = true) + val originalQuery = streamingRelation.where(10 < 5).select($"a").analyze + val optimized = Optimize.execute(originalQuery) + val correctAnswer = streamingRelation.where(10 < 5).select($"a").analyze + comparePlans(optimized, correctAnswer) + } + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala index d5facc245e72f..e1e5b3a7ef88e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/OffsetSeq.scala @@ -101,7 +101,9 @@ object OffsetSeqMetadata extends Logging { SHUFFLE_PARTITIONS, STATE_STORE_PROVIDER_CLASS, STREAMING_MULTIPLE_WATERMARK_POLICY, FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION, STREAMING_AGGREGATION_STATE_FORMAT_VERSION, STREAMING_JOIN_STATE_FORMAT_VERSION, STATE_STORE_COMPRESSION_CODEC, - STATE_STORE_ROCKSDB_FORMAT_VERSION, STATEFUL_OPERATOR_USE_STRICT_DISTRIBUTION) + STATE_STORE_ROCKSDB_FORMAT_VERSION, STATEFUL_OPERATOR_USE_STRICT_DISTRIBUTION, + PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN + ) /** * Default values of relevant configurations that are used for backward compatibility. @@ -122,7 +124,8 @@ object OffsetSeqMetadata extends Logging { STREAMING_JOIN_STATE_FORMAT_VERSION.key -> SymmetricHashJoinStateManager.legacyVersion.toString, STATE_STORE_COMPRESSION_CODEC.key -> CompressionCodec.LZ4, - STATEFUL_OPERATOR_USE_STRICT_DISTRIBUTION.key -> "false" + STATEFUL_OPERATOR_USE_STRICT_DISTRIBUTION.key -> "false", + PRUNE_FILTERS_CAN_PRUNE_STREAMING_SUBPLAN.key -> "true" ) def apply(json: String): OffsetSeqMetadata = Serialization.read[OffsetSeqMetadata](json) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryOptimizationCorrectnessSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryOptimizationCorrectnessSuite.scala index 782badaef924f..f651bfb7f3c72 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryOptimizationCorrectnessSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryOptimizationCorrectnessSuite.scala @@ -21,7 +21,7 @@ import java.sql.Timestamp import org.apache.spark.sql.Row import org.apache.spark.sql.execution.streaming.MemoryStream -import org.apache.spark.sql.functions.{expr, lit, window} +import org.apache.spark.sql.functions.{count, expr, lit, timestamp_seconds, window} import org.apache.spark.sql.internal.SQLConf /** @@ -524,4 +524,66 @@ class StreamingQueryOptimizationCorrectnessSuite extends StreamTest { doTest(numExpectedStatefulOperatorsForOneEmptySource = 1) } } + + test("SPARK-49699: observe node is not pruned out from PruneFilters") { + val input1 = MemoryStream[Int] + val df = input1.toDF() + .withColumn("eventTime", timestamp_seconds($"value")) + .observe("observation", count(lit(1)).as("rows")) + // Enforce PruneFilters to come into play and prune subtree. We could do the same + // with the reproducer of SPARK-48267, but let's just be simpler. + .filter(expr("false")) + + testStream(df)( + AddData(input1, 1, 2, 3), + CheckNewAnswer(), + Execute { qe => + val observeRow = qe.lastExecution.observedMetrics.get("observation") + assert(observeRow.get.getAs[Long]("rows") == 3L) + } + ) + } + + test("SPARK-49699: watermark node is not pruned out from PruneFilters") { + // NOTE: The test actually passes without SPARK-49699, because of the trickiness of + // filter pushdown and PruneFilters. Unlike observe node, the `false` filter is pushed down + // below to watermark node, hence PruneFilters rule does not prune out watermark node even + // before SPARK-49699. Propagate empty relation does not also propagate emptiness into + // watermark node, so the node is retained. The test is added for preventing regression. + + val input1 = MemoryStream[Int] + val df = input1.toDF() + .withColumn("eventTime", timestamp_seconds($"value")) + .withWatermark("eventTime", "0 second") + // Enforce PruneFilter to come into play and prune subtree. We could do the same + // with the reproducer of SPARK-48267, but let's just be simpler. + .filter(expr("false")) + + testStream(df)( + AddData(input1, 1, 2, 3), + CheckNewAnswer(), + Execute { qe => + // If the watermark node is pruned out, this would be null. + assert(qe.lastProgress.eventTime.get("watermark") != null) + } + ) + } + + test("SPARK-49699: stateful operator node is not pruned out from PruneFilters") { + val input1 = MemoryStream[Int] + val df = input1.toDF() + .groupBy("value") + .count() + // Enforce PruneFilter to come into play and prune subtree. We could do the same + // with the reproducer of SPARK-48267, but let's just be simpler. + .filter(expr("false")) + + testStream(df, OutputMode.Complete())( + AddData(input1, 1, 2, 3), + CheckNewAnswer(), + Execute { qe => + assert(qe.lastProgress.stateOperators.length == 1) + } + ) + } } From 4d97574425e603a7c6ac42a419747922bb1f83f9 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Fri, 20 Sep 2024 15:15:14 +0200 Subject: [PATCH 086/189] [SPARK-49733][SQL][DOCS] Delete `ExpressionInfo[between]` from `gen-sql-api-docs.py` to avoid duplication ### What changes were proposed in this pull request? The pr aims to delete `ExpressionInfo[between]` from `gen-sql-api-docs.py` to avoid duplication. ### Why are the changes needed? - In the following doc, `between` is repeatedly displayed `twice` https://spark.apache.org/docs/preview/api/sql/index.html#between image After the pr: image - After https://github.com/apache/spark/pull/44299, the expression 'between' has been added to `Spark 4.0`. ### Does this PR introduce _any_ user-facing change? Yes, only for docs. ### How was this patch tested? Manually check. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48183 from panbingkun/SPARK-49733. Authored-by: panbingkun Signed-off-by: Max Gekk --- .../spark/sql/catalyst/expressions/Between.scala | 2 +- sql/gen-sql-api-docs.py | 13 ------------- 2 files changed, 1 insertion(+), 14 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala index de1122da646b7..deec1ab51ad98 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.internal.SQLConf * lower - Lower bound of the between check. * upper - Upper bound of the between check. """, - since = "4.0.0", + since = "1.0.0", group = "conditional_funcs") case class Between private(input: Expression, lower: Expression, upper: Expression, replacement: Expression) extends RuntimeReplaceable with InheritAnalysisRules { diff --git a/sql/gen-sql-api-docs.py b/sql/gen-sql-api-docs.py index 17631a7352a02..3d19da01b3938 100644 --- a/sql/gen-sql-api-docs.py +++ b/sql/gen-sql-api-docs.py @@ -69,19 +69,6 @@ note="", since="1.0.0", deprecated=""), - ExpressionInfo( - className="", - name="between", - usage="expr1 [NOT] BETWEEN expr2 AND expr3 - " + - "evaluate if `expr1` is [not] in between `expr2` and `expr3`.", - arguments="", - examples="\n Examples:\n " + - "> SELECT col1 FROM VALUES 1, 3, 5, 7 WHERE col1 BETWEEN 2 AND 5;\n " + - " 3\n " + - " 5", - note="", - since="1.0.0", - deprecated=""), ExpressionInfo( className="", name="case", From bb8294c649909702e9086203b2726c6f51971c9c Mon Sep 17 00:00:00 2001 From: panbingkun Date: Fri, 20 Sep 2024 16:09:00 +0200 Subject: [PATCH 087/189] [SPARK-49729][SQL][DOCS] Forcefully check `usage` and correct the non-standard writing of 4 expressions ### What changes were proposed in this pull request? The pr aims to - forcefully check `usage` - correct the non-standard writing of 4 expressions (`shiftleft`, `shiftright`, `shiftrightunsigned`, `between`) ### Why are the changes needed? 1.When some expressions have non-standard `usage` writing, corresponding explanations may be omitted in our documentation, such as `shiftleft` https://spark.apache.org/docs/preview/sql-ref-functions-builtin.html - Before (Note: It looks very weird to only appear in `examples` and not in the `Conditional Functions` catalog) image - After image 2.When there is an `non-standard` writing format, it fails directly in GA and can be corrected in a timely manner to avoid omissions. Refer to `Manually check` below. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Pass GA. - Manually check: ```python The usage of between is not standardized, please correct it. Refer to: `AesDecrypt` ------------------------------------------------ Jekyll 4.3.3 Please append `--trace` to the `build` command for any additional information or backtrace. ------------------------------------------------ /Users/panbingkun/Developer/spark/spark-community/docs/_plugins/build_api_docs.rb:184:in `build_sql_docs': SQL doc generation failed (RuntimeError) from /Users/panbingkun/Developer/spark/spark-community/docs/_plugins/build_api_docs.rb:225:in `' from :37:in `require' from :37:in `require' from /Users/panbingkun/Developer/spark/spark-community/docs/.local_ruby_bundle/ruby/3.3.0/gems/jekyll-4.3.3/lib/jekyll/external.rb:57:in `block in require_with_graceful_fail' from /Users/panbingkun/Developer/spark/spark-community/docs/.local_ruby_bundle/ruby/3.3.0/gems/jekyll-4.3.3/lib/jekyll/external.rb:55:in `each' from /Users/panbingkun/Developer/spark/spark-community/docs/.local_ruby_bundle/ruby/3.3.0/gems/jekyll-4.3.3/lib/jekyll/external.rb:55:in `require_with_graceful_fail' from /Users/panbingkun/Developer/spark/spark-community/docs/.local_ruby_bundle/ruby/3.3.0/gems/jekyll-4.3.3/lib/jekyll/plugin_manager.rb:96:in `block in require_plugin_files' ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48179 from panbingkun/SPARK-49729. Authored-by: panbingkun Signed-off-by: Max Gekk --- .../spark/sql/catalyst/expressions/Between.scala | 2 +- .../sql/catalyst/expressions/mathExpressions.scala | 6 +++--- sql/gen-sql-functions-docs.py | 12 +++++++++++- .../spark/sql/hive/execution/SQLQuerySuite.scala | 2 +- 4 files changed, 16 insertions(+), 6 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala index deec1ab51ad98..c226e48c6be5e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Between.scala @@ -21,7 +21,7 @@ import org.apache.spark.sql.internal.SQLConf // scalastyle:off line.size.limit @ExpressionDescription( - usage = "Usage: input [NOT] BETWEEN lower AND upper - evaluate if `input` is [not] in between `lower` and `upper`", + usage = "input [NOT] _FUNC_ lower AND upper - evaluate if `input` is [not] in between `lower` and `upper`", examples = """ Examples: > SELECT 0.5 _FUNC_ 0.1 AND 1.0; diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala index 00274a16b888b..ddba820414ae4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala @@ -1293,7 +1293,7 @@ sealed trait BitShiftOperation * @param right number of bits to left shift. */ @ExpressionDescription( - usage = "base << exp - Bitwise left shift.", + usage = "base _FUNC_ exp - Bitwise left shift.", examples = """ Examples: > SELECT shiftleft(2, 1); @@ -1322,7 +1322,7 @@ case class ShiftLeft(left: Expression, right: Expression) extends BitShiftOperat * @param right number of bits to right shift. */ @ExpressionDescription( - usage = "base >> expr - Bitwise (signed) right shift.", + usage = "base _FUNC_ expr - Bitwise (signed) right shift.", examples = """ Examples: > SELECT shiftright(4, 1); @@ -1350,7 +1350,7 @@ case class ShiftRight(left: Expression, right: Expression) extends BitShiftOpera * @param right the number of bits to right shift. */ @ExpressionDescription( - usage = "base >>> expr - Bitwise unsigned right shift.", + usage = "base _FUNC_ expr - Bitwise unsigned right shift.", examples = """ Examples: > SELECT shiftrightunsigned(4, 1); diff --git a/sql/gen-sql-functions-docs.py b/sql/gen-sql-functions-docs.py index bb813cffb0128..4be9966747d1f 100644 --- a/sql/gen-sql-functions-docs.py +++ b/sql/gen-sql-functions-docs.py @@ -39,6 +39,10 @@ } +def _print_red(text): + print('\033[31m' + text + '\033[0m') + + def _list_grouped_function_infos(jvm): """ Returns a list of function information grouped by each group value via JVM. @@ -126,7 +130,13 @@ def _make_pretty_usage(infos): func_name = "\\" + func_name elif (info.name == "when"): func_name = "CASE WHEN" - usages = iter(re.split(r"(.*%s.*) - " % func_name, info.usage.strip())[1:]) + expr_usages = re.split(r"(.*%s.*) - " % func_name, info.usage.strip()) + if len(expr_usages) <= 1: + _print_red("\nThe `usage` of %s is not standardized, please correct it. " + "Refer to: `AesDecrypt`" % (func_name)) + os._exit(-1) + usages = iter(expr_usages[1:]) + for (sig, description) in zip(usages, usages): result.append(" ") result.append(" %s" % sig) diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala index 594c097de2c7d..14051034a588e 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala @@ -246,7 +246,7 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi checkKeywordsExist(sql("describe function `between`"), "Function: between", - "Usage: input [NOT] BETWEEN lower AND upper - " + + "input [NOT] between lower AND upper - " + "evaluate if `input` is [not] in between `lower` and `upper`") checkKeywordsExist(sql("describe function `case`"), From 3d8c078ddefe3bb74fc78ffc9391a067156c8499 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Fri, 20 Sep 2024 08:44:14 -0700 Subject: [PATCH 088/189] [SPARK-49704][BUILD] Upgrade `commons-io` to 2.17.0 ### What changes were proposed in this pull request? This PR aims to upgrade `commons-io` from `2.16.1` to `2.17.0`. ### Why are the changes needed? The full release notes: https://commons.apache.org/proper/commons-io/changes-report.html#a2.17.0 ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48154 from panbingkun/SPARK-49704. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 2 +- pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 9871cc0bca04f..419625f48fa11 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -44,7 +44,7 @@ commons-compiler/3.1.9//commons-compiler-3.1.9.jar commons-compress/1.27.1//commons-compress-1.27.1.jar commons-crypto/1.1.0//commons-crypto-1.1.0.jar commons-dbcp/1.4//commons-dbcp-1.4.jar -commons-io/2.16.1//commons-io-2.16.1.jar +commons-io/2.17.0//commons-io-2.17.0.jar commons-lang/2.6//commons-lang-2.6.jar commons-lang3/3.17.0//commons-lang3-3.17.0.jar commons-math3/3.6.1//commons-math3-3.6.1.jar diff --git a/pom.xml b/pom.xml index ddabc82d2ad13..b7c87beec0f92 100644 --- a/pom.xml +++ b/pom.xml @@ -187,7 +187,7 @@ 3.0.3 1.17.1 1.27.1 - 2.16.1 + 2.17.0 2.6 From 22a7edce0a7c70d6c1a5dcf995c6c723f0c3352b Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Fri, 20 Sep 2024 08:53:52 -0700 Subject: [PATCH 089/189] [SPARK-49531][PYTHON][CONNECT] Support line plot with plotly backend ### What changes were proposed in this pull request? Support line plot with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations, such as line plots, by leveraging libraries like Plotly. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. ```python >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] >>> columns = ["category", "int_val", "float_val"] >>> sdf = spark.createDataFrame(data, columns) >>> sdf.show() +--------+-------+---------+ |category|int_val|float_val| +--------+-------+---------+ | A| 10| 1.5| | B| 30| 2.5| | C| 20| 3.5| +--------+-------+---------+ >>> f = sdf.plot(kind="line", x="category", y="int_val") >>> f.show() # see below >>> g = sdf.plot.line(x="category", y=["int_val", "float_val"]) >>> g.show() # see below ``` `f.show()`: ![newplot](https://github.com/user-attachments/assets/ebd50bbc-0dd1-437f-ae0c-0b4de8f3c722) `g.show()`: ![newplot (1)](https://github.com/user-attachments/assets/46d28840-a147-428f-8d88-d424aa76ad06) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48139 from xinrong-meng/plot_line_w_dep. Authored-by: Xinrong Meng Signed-off-by: Dongjoon Hyun --- .github/workflows/build_python_connect.yml | 2 +- dev/requirements.txt | 2 +- dev/sparktestsupport/modules.py | 4 + .../docs/source/getting_started/install.rst | 1 + python/packaging/classic/setup.py | 1 + python/packaging/connect/setup.py | 2 + python/pyspark/errors/error-conditions.json | 5 + python/pyspark/sql/classic/dataframe.py | 9 ++ python/pyspark/sql/connect/dataframe.py | 8 ++ python/pyspark/sql/dataframe.py | 28 ++++ python/pyspark/sql/plot/__init__.py | 21 +++ python/pyspark/sql/plot/core.py | 135 ++++++++++++++++++ python/pyspark/sql/plot/plotly.py | 30 ++++ .../tests/connect/test_parity_frame_plot.py | 36 +++++ .../connect/test_parity_frame_plot_plotly.py | 36 +++++ python/pyspark/sql/tests/plot/__init__.py | 16 +++ .../pyspark/sql/tests/plot/test_frame_plot.py | 80 +++++++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 64 +++++++++ python/pyspark/sql/utils.py | 17 +++ python/pyspark/testing/sqlutils.py | 7 + .../apache/spark/sql/internal/SQLConf.scala | 27 ++++ 21 files changed, 529 insertions(+), 2 deletions(-) create mode 100644 python/pyspark/sql/plot/__init__.py create mode 100644 python/pyspark/sql/plot/core.py create mode 100644 python/pyspark/sql/plot/plotly.py create mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot.py create mode 100644 python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py create mode 100644 python/pyspark/sql/tests/plot/__init__.py create mode 100644 python/pyspark/sql/tests/plot/test_frame_plot.py create mode 100644 python/pyspark/sql/tests/plot/test_frame_plot_plotly.py diff --git a/.github/workflows/build_python_connect.yml b/.github/workflows/build_python_connect.yml index 3ac1a0117e41b..f668d813ef26e 100644 --- a/.github/workflows/build_python_connect.yml +++ b/.github/workflows/build_python_connect.yml @@ -71,7 +71,7 @@ jobs: python packaging/connect/setup.py sdist cd dist pip install pyspark*connect-*.tar.gz - pip install 'six==1.16.0' 'pandas<=2.2.2' scipy 'plotly>=4.8' 'mlflow>=2.8.1' coverage matplotlib openpyxl 'memory-profiler>=0.61.0' 'scikit-learn>=1.3.2' 'graphviz==0.20.3' torch torchvision torcheval deepspeed unittest-xml-reporting + pip install 'six==1.16.0' 'pandas<=2.2.2' scipy 'plotly>=4.8' 'mlflow>=2.8.1' coverage matplotlib openpyxl 'memory-profiler>=0.61.0' 'scikit-learn>=1.3.2' 'graphviz==0.20.3' torch torchvision torcheval deepspeed unittest-xml-reporting 'plotly>=4.8' - name: Run tests env: SPARK_TESTING: 1 diff --git a/dev/requirements.txt b/dev/requirements.txt index 5486c98ab8f8f..cafc73405aaa8 100644 --- a/dev/requirements.txt +++ b/dev/requirements.txt @@ -7,7 +7,7 @@ pyarrow>=10.0.0 six==1.16.0 pandas>=2.0.0 scipy -plotly +plotly>=4.8 mlflow>=2.3.1 scikit-learn matplotlib diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 34fbb8450d544..b9a4bed715f67 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -548,6 +548,8 @@ def __hash__(self): "pyspark.sql.tests.test_udtf", "pyspark.sql.tests.test_utils", "pyspark.sql.tests.test_resources", + "pyspark.sql.tests.plot.test_frame_plot", + "pyspark.sql.tests.plot.test_frame_plot_plotly", ], ) @@ -1051,6 +1053,8 @@ def __hash__(self): "pyspark.sql.tests.connect.test_parity_arrow_cogrouped_map", "pyspark.sql.tests.connect.test_parity_python_datasource", "pyspark.sql.tests.connect.test_parity_python_streaming_datasource", + "pyspark.sql.tests.connect.test_parity_frame_plot", + "pyspark.sql.tests.connect.test_parity_frame_plot_plotly", "pyspark.sql.tests.connect.test_utils", "pyspark.sql.tests.connect.client.test_artifact", "pyspark.sql.tests.connect.client.test_artifact_localcluster", diff --git a/python/docs/source/getting_started/install.rst b/python/docs/source/getting_started/install.rst index 549656bea103e..88c0a8c26cc94 100644 --- a/python/docs/source/getting_started/install.rst +++ b/python/docs/source/getting_started/install.rst @@ -183,6 +183,7 @@ Package Supported version Note Additional libraries that enhance functionality but are not included in the installation packages: - **memory-profiler**: Used for PySpark UDF memory profiling, ``spark.profile.show(...)`` and ``spark.sql.pyspark.udf.profiler``. +- **plotly**: Used for PySpark plotting, ``DataFrame.plot``. Note that PySpark requires Java 17 or later with ``JAVA_HOME`` properly set and refer to |downloading|_. diff --git a/python/packaging/classic/setup.py b/python/packaging/classic/setup.py index 79b74483f00dd..17cca326d0241 100755 --- a/python/packaging/classic/setup.py +++ b/python/packaging/classic/setup.py @@ -288,6 +288,7 @@ def run(self): "pyspark.sql.connect.streaming.worker", "pyspark.sql.functions", "pyspark.sql.pandas", + "pyspark.sql.plot", "pyspark.sql.protobuf", "pyspark.sql.streaming", "pyspark.sql.worker", diff --git a/python/packaging/connect/setup.py b/python/packaging/connect/setup.py index ab166c79747df..6ae16e9a9ad3a 100755 --- a/python/packaging/connect/setup.py +++ b/python/packaging/connect/setup.py @@ -77,6 +77,7 @@ "pyspark.sql.tests.connect.client", "pyspark.sql.tests.connect.shell", "pyspark.sql.tests.pandas", + "pyspark.sql.tests.plot", "pyspark.sql.tests.streaming", "pyspark.ml.tests.connect", "pyspark.pandas.tests", @@ -161,6 +162,7 @@ "pyspark.sql.connect.streaming.worker", "pyspark.sql.functions", "pyspark.sql.pandas", + "pyspark.sql.plot", "pyspark.sql.protobuf", "pyspark.sql.streaming", "pyspark.sql.worker", diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 4061d024a83cd..92aeb15e21d1b 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -1088,6 +1088,11 @@ "Function `` should use only POSITIONAL or POSITIONAL OR KEYWORD arguments." ] }, + "UNSUPPORTED_PLOT_BACKEND": { + "message": [ + "`` is not supported, it should be one of the values from " + ] + }, "UNSUPPORTED_SIGNATURE": { "message": [ "Unsupported signature: ." diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index 91b9591625904..a2778cbc32c4c 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -73,6 +73,11 @@ from pyspark.sql.pandas.conversion import PandasConversionMixin from pyspark.sql.pandas.map_ops import PandasMapOpsMixin +try: + from pyspark.sql.plot import PySparkPlotAccessor +except ImportError: + PySparkPlotAccessor = None # type: ignore + if TYPE_CHECKING: from py4j.java_gateway import JavaObject import pyarrow as pa @@ -1862,6 +1867,10 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: messageParameters={"member": "queryExecution"}, ) + @property + def plot(self) -> PySparkPlotAccessor: + return PySparkPlotAccessor(self) + class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 768abd655d497..59d79decf6690 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -86,6 +86,10 @@ from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema from pyspark.sql.pandas.functions import _validate_pandas_udf # type: ignore[attr-defined] +try: + from pyspark.sql.plot import PySparkPlotAccessor +except ImportError: + PySparkPlotAccessor = None # type: ignore if TYPE_CHECKING: from pyspark.sql.connect._typing import ( @@ -2239,6 +2243,10 @@ def rdd(self) -> "RDD[Row]": def executionInfo(self) -> Optional["ExecutionInfo"]: return self._execution_info + @property + def plot(self) -> PySparkPlotAccessor: + return PySparkPlotAccessor(self) + class DataFrameNaFunctions(ParentDataFrameNaFunctions): def __init__(self, df: ParentDataFrame): diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index ef35b73332572..2179a844b1e5e 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -43,6 +43,7 @@ from pyspark.sql.types import StructType, Row from pyspark.sql.utils import dispatch_df_method + if TYPE_CHECKING: from py4j.java_gateway import JavaObject import pyarrow as pa @@ -65,6 +66,7 @@ ArrowMapIterFunction, DataFrameLike as PandasDataFrameLike, ) + from pyspark.sql.plot import PySparkPlotAccessor from pyspark.sql.metrics import ExecutionInfo @@ -6394,6 +6396,32 @@ def executionInfo(self) -> Optional["ExecutionInfo"]: """ ... + @property + def plot(self) -> "PySparkPlotAccessor": + """ + Returns a :class:`PySparkPlotAccessor` for plotting functions. + + .. versionadded:: 4.0.0 + + Returns + ------- + :class:`PySparkPlotAccessor` + + Notes + ----- + This API is experimental. + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> type(df.plot) + + >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP + """ + ... + class DataFrameNaFunctions: """Functionality for working with missing data in :class:`DataFrame`. diff --git a/python/pyspark/sql/plot/__init__.py b/python/pyspark/sql/plot/__init__.py new file mode 100644 index 0000000000000..6da07061b2a09 --- /dev/null +++ b/python/pyspark/sql/plot/__init__.py @@ -0,0 +1,21 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +This package includes the plotting APIs for PySpark DataFrame. +""" +from pyspark.sql.plot.core import * # noqa: F403, F401 diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py new file mode 100644 index 0000000000000..392ef73b38845 --- /dev/null +++ b/python/pyspark/sql/plot/core.py @@ -0,0 +1,135 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Any, TYPE_CHECKING, Optional, Union +from types import ModuleType +from pyspark.errors import PySparkRuntimeError, PySparkValueError +from pyspark.sql.utils import require_minimum_plotly_version + + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + import pandas as pd + from plotly.graph_objs import Figure + + +class PySparkTopNPlotBase: + def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": + from pyspark.sql import SparkSession + + session = SparkSession.getActiveSession() + if session is None: + raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) + + max_rows = int( + session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] + ) + pdf = sdf.limit(max_rows + 1).toPandas() + + self.partial = False + if len(pdf) > max_rows: + self.partial = True + pdf = pdf.iloc[:max_rows] + + return pdf + + +class PySparkSampledPlotBase: + def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": + from pyspark.sql import SparkSession + + session = SparkSession.getActiveSession() + if session is None: + raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) + + sample_ratio = session.conf.get("spark.sql.pyspark.plotting.sample_ratio") + max_rows = int( + session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] + ) + + if sample_ratio is None: + fraction = 1 / (sdf.count() / max_rows) + fraction = min(1.0, fraction) + else: + fraction = float(sample_ratio) + + sampled_sdf = sdf.sample(fraction=fraction) + pdf = sampled_sdf.toPandas() + + return pdf + + +class PySparkPlotAccessor: + plot_data_map = { + "line": PySparkSampledPlotBase().get_sampled, + } + _backends = {} # type: ignore[var-annotated] + + def __init__(self, data: "DataFrame"): + self.data = data + + def __call__( + self, kind: str = "line", backend: Optional[str] = None, **kwargs: Any + ) -> "Figure": + plot_backend = PySparkPlotAccessor._get_plot_backend(backend) + + return plot_backend.plot_pyspark(self.data, kind=kind, **kwargs) + + @staticmethod + def _get_plot_backend(backend: Optional[str] = None) -> ModuleType: + backend = backend or "plotly" + + if backend in PySparkPlotAccessor._backends: + return PySparkPlotAccessor._backends[backend] + + if backend == "plotly": + require_minimum_plotly_version() + else: + raise PySparkValueError( + errorClass="UNSUPPORTED_PLOT_BACKEND", + messageParameters={"backend": backend, "supported_backends": ", ".join(["plotly"])}, + ) + from pyspark.sql.plot import plotly as module + + return module + + def line(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": + """ + Plot DataFrame as lines. + + Parameters + ---------- + x : str + Name of column to use for the horizontal axis. + y : str or list of str + Name(s) of the column(s) to use for the vertical axis. Multiple columns can be plotted. + **kwargs : optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.line(x="category", y="int_val") # doctest: +SKIP + >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP + """ + return self(kind="line", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/plot/plotly.py b/python/pyspark/sql/plot/plotly.py new file mode 100644 index 0000000000000..5efc19476057f --- /dev/null +++ b/python/pyspark/sql/plot/plotly.py @@ -0,0 +1,30 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import TYPE_CHECKING, Any + +from pyspark.sql.plot import PySparkPlotAccessor + +if TYPE_CHECKING: + from pyspark.sql import DataFrame + from plotly.graph_objs import Figure + + +def plot_pyspark(data: "DataFrame", kind: str, **kwargs: Any) -> "Figure": + import plotly + + return plotly.plot(PySparkPlotAccessor.plot_data_map[kind](data), kind, **kwargs) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py new file mode 100644 index 0000000000000..c69e438bf7eb0 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py @@ -0,0 +1,36 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.sql.tests.plot.test_frame_plot import DataFramePlotTestsMixin + + +class FramePlotParityTests(DataFramePlotTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_frame_plot import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py new file mode 100644 index 0000000000000..78508fe533379 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py @@ -0,0 +1,36 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.sql.tests.plot.test_frame_plot_plotly import DataFramePlotPlotlyTestsMixin + + +class FramePlotPlotlyParityTests(DataFramePlotPlotlyTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + import unittest + from pyspark.sql.tests.connect.test_parity_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/__init__.py b/python/pyspark/sql/tests/plot/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/sql/tests/plot/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/sql/tests/plot/test_frame_plot.py b/python/pyspark/sql/tests/plot/test_frame_plot.py new file mode 100644 index 0000000000000..f753b5ab3db72 --- /dev/null +++ b/python/pyspark/sql/tests/plot/test_frame_plot.py @@ -0,0 +1,80 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +from pyspark.errors import PySparkValueError +from pyspark.sql import Row +from pyspark.sql.plot import PySparkSampledPlotBase, PySparkTopNPlotBase +from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message + + +@unittest.skipIf(not have_plotly, plotly_requirement_message) +class DataFramePlotTestsMixin: + def test_backend(self): + accessor = self.spark.range(2).plot + backend = accessor._get_plot_backend() + self.assertEqual(backend.__name__, "pyspark.sql.plot.plotly") + + with self.assertRaises(PySparkValueError) as pe: + accessor._get_plot_backend("matplotlib") + + self.check_error( + exception=pe.exception, + errorClass="UNSUPPORTED_PLOT_BACKEND", + messageParameters={"backend": "matplotlib", "supported_backends": "plotly"}, + ) + + def test_topn_max_rows(self): + try: + self.spark.conf.set("spark.sql.pyspark.plotting.max_rows", "1000") + sdf = self.spark.range(2500) + pdf = PySparkTopNPlotBase().get_top_n(sdf) + self.assertEqual(len(pdf), 1000) + finally: + self.spark.conf.unset("spark.sql.pyspark.plotting.max_rows") + + def test_sampled_plot_with_ratio(self): + try: + self.spark.conf.set("spark.sql.pyspark.plotting.sample_ratio", "0.5") + data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2500)] + sdf = self.spark.createDataFrame(data) + pdf = PySparkSampledPlotBase().get_sampled(sdf) + self.assertEqual(round(len(pdf) / 2500, 1), 0.5) + finally: + self.spark.conf.unset("spark.sql.pyspark.plotting.sample_ratio") + + def test_sampled_plot_with_max_rows(self): + data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2000)] + sdf = self.spark.createDataFrame(data) + pdf = PySparkSampledPlotBase().get_sampled(sdf) + self.assertEqual(round(len(pdf) / 2000, 1), 0.5) + + +class DataFramePlotTests(DataFramePlotTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.plot.test_frame_plot import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py new file mode 100644 index 0000000000000..72a3ed267d192 --- /dev/null +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -0,0 +1,64 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +import pyspark.sql.plot # noqa: F401 +from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message + + +@unittest.skipIf(not have_plotly, plotly_requirement_message) +class DataFramePlotPlotlyTestsMixin: + @property + def sdf(self): + data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + columns = ["category", "int_val", "float_val"] + return self.spark.createDataFrame(data, columns) + + def _check_fig_data(self, fig_data, expected_x, expected_y, expected_name=""): + self.assertEqual(fig_data["mode"], "lines") + self.assertEqual(fig_data["type"], "scatter") + self.assertEqual(fig_data["xaxis"], "x") + self.assertEqual(list(fig_data["x"]), expected_x) + self.assertEqual(fig_data["yaxis"], "y") + self.assertEqual(list(fig_data["y"]), expected_y) + self.assertEqual(fig_data["name"], expected_name) + + def test_line_plot(self): + # single column as vertical axis + fig = self.sdf.plot(kind="line", x="category", y="int_val") + self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20]) + + # multiple columns as vertical axis + fig = self.sdf.plot.line(x="category", y=["int_val", "float_val"]) + self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") + self._check_fig_data(fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + + +class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.plot.test_frame_plot_plotly import * # noqa: F401 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py index 11b91612419a3..5d9ec92cbc830 100644 --- a/python/pyspark/sql/utils.py +++ b/python/pyspark/sql/utils.py @@ -41,6 +41,7 @@ PythonException, UnknownException, SparkUpgradeException, + PySparkImportError, PySparkNotImplementedError, PySparkRuntimeError, ) @@ -115,6 +116,22 @@ def require_test_compiled() -> None: ) +def require_minimum_plotly_version() -> None: + """Raise ImportError if plotly is not installed""" + minimum_plotly_version = "4.8" + + try: + import plotly # noqa: F401 + except ImportError as error: + raise PySparkImportError( + errorClass="PACKAGE_NOT_INSTALLED", + messageParameters={ + "package_name": "plotly", + "minimum_version": str(minimum_plotly_version), + }, + ) from error + + class ForeachBatchFunction: """ This is the Python implementation of Java interface 'ForeachBatchFunction'. This wraps diff --git a/python/pyspark/testing/sqlutils.py b/python/pyspark/testing/sqlutils.py index 9f07c44c084cf..00ad40e68bd7c 100644 --- a/python/pyspark/testing/sqlutils.py +++ b/python/pyspark/testing/sqlutils.py @@ -48,6 +48,13 @@ except Exception as e: test_not_compiled_message = str(e) +plotly_requirement_message = None +try: + import plotly +except ImportError as e: + plotly_requirement_message = str(e) +have_plotly = plotly_requirement_message is None + from pyspark.sql import SparkSession from pyspark.sql.types import ArrayType, DoubleType, UserDefinedType, Row from pyspark.testing.utils import ReusedPySparkTestCase, PySparkErrorTestUtils diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 2eaafde52228b..6c3e9bac1cfe5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3169,6 +3169,29 @@ object SQLConf { .version("4.0.0") .fallbackConf(Python.PYTHON_WORKER_FAULTHANLDER_ENABLED) + val PYSPARK_PLOT_MAX_ROWS = + buildConf("spark.sql.pyspark.plotting.max_rows") + .doc( + "The visual limit on top-n-based plots. If set to 1000, the first 1000 data points " + + "will be used for plotting.") + .version("4.0.0") + .intConf + .createWithDefault(1000) + + val PYSPARK_PLOT_SAMPLE_RATIO = + buildConf("spark.sql.pyspark.plotting.sample_ratio") + .doc( + "The proportion of data that will be plotted for sample-based plots. It is determined " + + "based on spark.sql.pyspark.plotting.max_rows if not explicitly set." + ) + .version("4.0.0") + .doubleConf + .checkValue( + ratio => ratio >= 0.0 && ratio <= 1.0, + "The value should be between 0.0 and 1.0 inclusive." + ) + .createOptional + val ARROW_SPARKR_EXECUTION_ENABLED = buildConf("spark.sql.execution.arrow.sparkr.enabled") .doc("When true, make use of Apache Arrow for columnar data transfers in SparkR. " + @@ -5873,6 +5896,10 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pythonUDFWorkerFaulthandlerEnabled: Boolean = getConf(PYTHON_UDF_WORKER_FAULTHANLDER_ENABLED) + def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS) + + def pysparkPlotSampleRatio: Option[Double] = getConf(PYSPARK_PLOT_SAMPLE_RATIO) + def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED) def arrowPySparkFallbackEnabled: Boolean = getConf(ARROW_PYSPARK_FALLBACK_ENABLED) From f3785fadec3089fa60d85fa3c98ae9c6ada807a4 Mon Sep 17 00:00:00 2001 From: Stefan Kandic Date: Fri, 20 Sep 2024 19:12:05 +0200 Subject: [PATCH 090/189] [SPARK-49737][SQL] Disable bucketing on collated columns in complex types ### What changes were proposed in this pull request? To disable bucketing on collated string types in complex types (structs, arrays and maps). ### Why are the changes needed? #45260 introduces the logic to disabled bucketing for collated columns, but forgot to address complex types which have collated strings inside. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48186 from stefankandic/fixBucketing. Authored-by: Stefan Kandic Signed-off-by: Max Gekk --- .../datasources/BucketingUtils.scala | 8 +++---- .../org/apache/spark/sql/CollationSuite.scala | 23 ++++++++++++++----- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/BucketingUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/BucketingUtils.scala index 4fa1e0c1f2c58..fd47feef25d57 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/BucketingUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/BucketingUtils.scala @@ -19,7 +19,8 @@ package org.apache.spark.sql.execution.datasources import org.apache.spark.sql.catalyst.expressions.{Attribute, SpecificInternalRow, UnsafeProjection} import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning -import org.apache.spark.sql.types.{DataType, StringType} +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.util.SchemaUtils object BucketingUtils { // The file name of bucketed data should have 3 parts: @@ -53,10 +54,7 @@ object BucketingUtils { bucketIdGenerator(mutableInternalRow).getInt(0) } - def canBucketOn(dataType: DataType): Boolean = dataType match { - case st: StringType => st.supportsBinaryOrdering - case other => true - } + def canBucketOn(dataType: DataType): Boolean = !SchemaUtils.hasNonUTF8BinaryCollation(dataType) def bucketIdToString(id: Int): String = f"_$id%05d" } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index 73fd897e91f53..632b9305feb57 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -162,9 +162,14 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { withTable(tableName) { sql( s""" - |CREATE TABLE $tableName - |(id INT, c1 STRING COLLATE UNICODE, c2 string) - |USING parquet + |CREATE TABLE $tableName ( + | id INT, + | c1 STRING COLLATE UNICODE, + | c2 STRING, + | struct_col STRUCT, + | array_col ARRAY, + | map_col MAP + |) USING parquet |CLUSTERED BY (${bucketColumns.mkString(",")}) |INTO 4 BUCKETS""".stripMargin ) @@ -175,14 +180,20 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { createTable("c2") createTable("id", "c2") - Seq(Seq("c1"), Seq("c1", "id"), Seq("c1", "c2")).foreach { bucketColumns => + val failBucketingColumns = Seq( + Seq("c1"), Seq("c1", "id"), Seq("c1", "c2"), + Seq("struct_col"), Seq("array_col"), Seq("map_col") + ) + + failBucketingColumns.foreach { bucketColumns => checkError( exception = intercept[AnalysisException] { createTable(bucketColumns: _*) }, condition = "INVALID_BUCKET_COLUMN_DATA_TYPE", - parameters = Map("type" -> "\"STRING COLLATE UNICODE\"") - ); + parameters = Map("type" -> ".*STRING COLLATE UNICODE.*"), + matchPVals = true + ) } } From f76a9b1135e748649bdb9a2104360f0dc533cc1f Mon Sep 17 00:00:00 2001 From: viktorluc-db Date: Fri, 20 Sep 2024 22:47:30 +0200 Subject: [PATCH 091/189] [SPARK-49738][SQL] Endswith bug fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Bugfix in "endswith" string predicate. Also fixed the same type of the bug in `CollationAwareUTF8String.java` in method `lowercaseMatchLengthFrom`. ### Why are the changes needed? Expression `select endswith('İo' collate utf8_lcase, 'İo' collate utf8_lcase)` returns `false` but should return `true`. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Added tests in CollationSupportSuite. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48187 from viktorluc-db/matchBugFix. Authored-by: viktorluc-db Signed-off-by: Max Gekk --- .../spark/sql/catalyst/util/CollationAwareUTF8String.java | 4 ++-- .../org/apache/spark/unsafe/types/CollationSupportSuite.java | 4 ++++ .../src/test/resources/sql-tests/results/collations.sql.out | 4 ++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationAwareUTF8String.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationAwareUTF8String.java index 5ed3048fb72b3..fb610a5d96f17 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationAwareUTF8String.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationAwareUTF8String.java @@ -109,7 +109,7 @@ private static int lowercaseMatchLengthFrom( } // Compare the characters in the target and pattern strings. int matchLength = 0, codePointBuffer = -1, targetCodePoint, patternCodePoint; - while (targetIterator.hasNext() && patternIterator.hasNext()) { + while ((targetIterator.hasNext() || codePointBuffer != -1) && patternIterator.hasNext()) { if (codePointBuffer != -1) { targetCodePoint = codePointBuffer; codePointBuffer = -1; @@ -211,7 +211,7 @@ private static int lowercaseMatchLengthUntil( } // Compare the characters in the target and pattern strings. int matchLength = 0, codePointBuffer = -1, targetCodePoint, patternCodePoint; - while (targetIterator.hasNext() && patternIterator.hasNext()) { + while ((targetIterator.hasNext() || codePointBuffer != -1) && patternIterator.hasNext()) { if (codePointBuffer != -1) { targetCodePoint = codePointBuffer; codePointBuffer = -1; diff --git a/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java b/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java index 5719303a0dce8..a445cde52ad57 100644 --- a/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java +++ b/common/unsafe/src/test/java/org/apache/spark/unsafe/types/CollationSupportSuite.java @@ -629,6 +629,8 @@ public void testStartsWith() throws SparkException { assertStartsWith("İonic", "Io", "UTF8_LCASE", false); assertStartsWith("İonic", "i\u0307o", "UTF8_LCASE", true); assertStartsWith("İonic", "İo", "UTF8_LCASE", true); + assertStartsWith("oİ", "oİ", "UTF8_LCASE", true); + assertStartsWith("oİ", "oi̇", "UTF8_LCASE", true); // Conditional case mapping (e.g. Greek sigmas). assertStartsWith("σ", "σ", "UTF8_BINARY", true); assertStartsWith("σ", "ς", "UTF8_BINARY", false); @@ -880,6 +882,8 @@ public void testEndsWith() throws SparkException { assertEndsWith("the İo", "Io", "UTF8_LCASE", false); assertEndsWith("the İo", "i\u0307o", "UTF8_LCASE", true); assertEndsWith("the İo", "İo", "UTF8_LCASE", true); + assertEndsWith("İo", "İo", "UTF8_LCASE", true); + assertEndsWith("İo", "i̇o", "UTF8_LCASE", true); // Conditional case mapping (e.g. Greek sigmas). assertEndsWith("σ", "σ", "UTF8_BINARY", true); assertEndsWith("σ", "ς", "UTF8_BINARY", false); diff --git a/sql/core/src/test/resources/sql-tests/results/collations.sql.out b/sql/core/src/test/resources/sql-tests/results/collations.sql.out index 5999bf20f6884..9d29a46e5a0ef 100644 --- a/sql/core/src/test/resources/sql-tests/results/collations.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/collations.sql.out @@ -2213,8 +2213,8 @@ struct Date: Fri, 20 Sep 2024 15:34:17 -0700 Subject: [PATCH 092/189] [SPARK-49557][SQL] Add SQL pipe syntax for the WHERE operator ### What changes were proposed in this pull request? This PR adds SQL pipe syntax support for the WHERE operator. For example: ``` CREATE TABLE t(x INT, y STRING) USING CSV; INSERT INTO t VALUES (0, 'abc'), (1, 'def'); CREATE TABLE other(a INT, b INT) USING JSON; INSERT INTO other VALUES (1, 1), (1, 2), (2, 4); TABLE t |> WHERE x + LENGTH(y) < 4; 0 abc TABLE t |> WHERE (SELECT ANY_VALUE(a) FROM other WHERE x = a LIMIT 1) = 1 1 def TABLE t |> WHERE SUM(x) = 1 Error: aggregate functions are not allowed in the pipe operator |> WHERE clause ``` ### Why are the changes needed? The SQL pipe operator syntax will let users compose queries in a more flexible fashion. ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds a few unit test cases, but mostly relies on golden file test coverage. I did this to make sure the answers are correct as this feature is implemented and also so we can look at the analyzer output plans to ensure they look right as well. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48091 from dtenedor/pipe-where. Authored-by: Daniel Tenedorio Signed-off-by: Gengliang Wang --- .../sql/catalyst/parser/SqlBaseParser.g4 | 1 + .../sql/catalyst/parser/AstBuilder.scala | 15 +- .../analyzer-results/pipe-operators.sql.out | 272 ++++++++++++++++++ .../sql-tests/inputs/pipe-operators.sql | 94 +++++- .../sql-tests/results/pipe-operators.sql.out | 268 +++++++++++++++++ .../sql/execution/SparkSqlParserSuite.scala | 12 +- 6 files changed, 658 insertions(+), 4 deletions(-) diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index e591a43b84d1a..094f7f5315b80 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -1492,6 +1492,7 @@ version operatorPipeRightSide : selectClause + | whereClause ; // When `SQL_standard_keyword_behavior=true`, there are 2 kinds of keywords in Spark SQL. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index 52529bb4b789b..674005caaf1b2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -5876,7 +5876,20 @@ class AstBuilder extends DataTypeAstBuilder windowClause = null, relation = left, isPipeOperatorSelect = true) - }.get + }.getOrElse(Option(ctx.whereClause).map { c => + // Add a table subquery boundary between the new filter and the input plan if one does not + // already exist. This helps the analyzer behave as if we had added the WHERE clause after a + // table subquery containing the input plan. + val withSubqueryAlias = left match { + case s: SubqueryAlias => + s + case u: UnresolvedRelation => + u + case _ => + SubqueryAlias(SubqueryAlias.generateSubqueryName(), left) + } + withWhereClause(c, withSubqueryAlias) + }.get) } /** diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out index ab0635fef048b..c44ce153a2f41 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out @@ -255,6 +255,55 @@ Distinct +- Relation spark_catalog.default.t[x#x,y#x] csv +-- !query +table t +|> select * +-- !query analysis +Project [x#x, y#x] ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select * except (y) +-- !query analysis +Project [x#x] ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select /*+ repartition(3) */ * +-- !query analysis +Repartition 3, true ++- Project [x#x, y#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select /*+ repartition(3) */ distinct x +-- !query analysis +Repartition 3, true ++- Distinct + +- Project [x#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> select /*+ repartition(3) */ all x +-- !query analysis +Repartition 3, true ++- Project [x#x] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + -- !query table t |> select sum(x) as result @@ -297,6 +346,229 @@ org.apache.spark.sql.AnalysisException } +-- !query +table t +|> where true +-- !query analysis +Filter true ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where x + length(y) < 4 +-- !query analysis +Filter ((x#x + length(y#x)) < 4) ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where x + length(y) < 4 +|> where x + length(y) < 3 +-- !query analysis +Filter ((x#x + length(y#x)) < 3) ++- SubqueryAlias __auto_generated_subquery_name + +- Filter ((x#x + length(y#x)) < 4) + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +(select x, sum(length(y)) as sum_len from t group by x) +|> where x = 1 +-- !query analysis +Filter (x#x = 1) ++- SubqueryAlias __auto_generated_subquery_name + +- Aggregate [x#x], [x#x, sum(length(y#x)) AS sum_len#xL] + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where t.x = 1 +-- !query analysis +Filter (x#x = 1) ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where spark_catalog.default.t.x = 1 +-- !query analysis +Filter (x#x = 1) ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +(select col from st) +|> where col.i1 = 1 +-- !query analysis +Filter (col#x.i1 = 1) ++- SubqueryAlias __auto_generated_subquery_name + +- Project [col#x] + +- SubqueryAlias spark_catalog.default.st + +- Relation spark_catalog.default.st[x#x,col#x] parquet + + +-- !query +table st +|> where st.col.i1 = 2 +-- !query analysis +Filter (col#x.i1 = 2) ++- SubqueryAlias spark_catalog.default.st + +- Relation spark_catalog.default.st[x#x,col#x] parquet + + +-- !query +table t +|> where exists (select a from other where x = a limit 1) +-- !query analysis +Filter exists#x [x#x] +: +- GlobalLimit 1 +: +- LocalLimit 1 +: +- Project [a#x] +: +- Filter (outer(x#x) = a#x) +: +- SubqueryAlias spark_catalog.default.other +: +- Relation spark_catalog.default.other[a#x,b#x] json ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where (select any_value(a) from other where x = a limit 1) = 1 +-- !query analysis +Filter (scalar-subquery#x [x#x] = 1) +: +- GlobalLimit 1 +: +- LocalLimit 1 +: +- Aggregate [any_value(a#x, false) AS any_value(a)#x] +: +- Filter (outer(x#x) = a#x) +: +- SubqueryAlias spark_catalog.default.other +: +- Relation spark_catalog.default.other[a#x,b#x] json ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> where sum(x) = 1 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_WHERE_CONDITION", + "sqlState" : "42903", + "messageParameters" : { + "condition" : "\"(sum(x) = 1)\"", + "expressionList" : "sum(spark_catalog.default.t.x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 27, + "fragment" : "table t\n|> where sum(x) = 1" + } ] +} + + +-- !query +table t +|> where y = 'abc' or length(y) + sum(x) = 1 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_WHERE_CONDITION", + "sqlState" : "42903", + "messageParameters" : { + "condition" : "\"((y = abc) OR ((length(y) + sum(x)) = 1))\"", + "expressionList" : "sum(spark_catalog.default.t.x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 52, + "fragment" : "table t\n|> where y = 'abc' or length(y) + sum(x) = 1" + } ] +} + + +-- !query +table t +|> where first_value(x) over (partition by y) = 1 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_1034", + "messageParameters" : { + "clauseName" : "WHERE" + } +} + + +-- !query +select * from t where first_value(x) over (partition by y) = 1 +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_1034", + "messageParameters" : { + "clauseName" : "WHERE" + } +} + + +-- !query +table t +|> select x, length(y) as z +|> where x + length(y) < 4 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`y`", + "proposal" : "`x`, `z`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 57, + "stopIndex" : 57, + "fragment" : "y" + } ] +} + + +-- !query +(select x, sum(length(y)) as sum_len from t group by x) +|> where sum(length(y)) = 3 +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`y`", + "proposal" : "`x`, `sum_len`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 77, + "stopIndex" : 77, + "fragment" : "y" + } ] +} + + -- !query drop table t -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql index 7d0966e7f2095..49a72137ee047 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql @@ -12,7 +12,7 @@ drop table if exists st; create table st(x int, col struct) using parquet; insert into st values (1, (2, 3)); --- Selection operators: positive tests. +-- SELECT operators: positive tests. --------------------------------------- -- Selecting a constant. @@ -85,7 +85,24 @@ table t table t |> select distinct x, y; --- Selection operators: negative tests. +-- SELECT * is supported. +table t +|> select *; + +table t +|> select * except (y); + +-- Hints are supported. +table t +|> select /*+ repartition(3) */ *; + +table t +|> select /*+ repartition(3) */ distinct x; + +table t +|> select /*+ repartition(3) */ all x; + +-- SELECT operators: negative tests. --------------------------------------- -- Aggregate functions are not allowed in the pipe operator SELECT list. @@ -95,6 +112,79 @@ table t table t |> select y, length(y) + sum(x) as result; +-- WHERE operators: positive tests. +----------------------------------- + +-- Filtering with a constant predicate. +table t +|> where true; + +-- Filtering with a predicate based on attributes from the input relation. +table t +|> where x + length(y) < 4; + +-- Two consecutive filters are allowed. +table t +|> where x + length(y) < 4 +|> where x + length(y) < 3; + +-- It is possible to use the WHERE operator instead of the HAVING clause when processing the result +-- of aggregations. For example, this WHERE operator is equivalent to the normal SQL "HAVING x = 1". +(select x, sum(length(y)) as sum_len from t group by x) +|> where x = 1; + +-- Filtering by referring to the table or table subquery alias. +table t +|> where t.x = 1; + +table t +|> where spark_catalog.default.t.x = 1; + +-- Filtering using struct fields. +(select col from st) +|> where col.i1 = 1; + +table st +|> where st.col.i1 = 2; + +-- Expression subqueries in the WHERE clause. +table t +|> where exists (select a from other where x = a limit 1); + +-- Aggregations are allowed within expression subqueries in the pipe operator WHERE clause as long +-- no aggregate functions exist in the top-level expression predicate. +table t +|> where (select any_value(a) from other where x = a limit 1) = 1; + +-- WHERE operators: negative tests. +----------------------------------- + +-- Aggregate functions are not allowed in the top-level WHERE predicate. +-- (Note: to implement this behavior, perform the aggregation first separately and then add a +-- pipe-operator WHERE clause referring to the result of aggregate expression(s) therein). +table t +|> where sum(x) = 1; + +table t +|> where y = 'abc' or length(y) + sum(x) = 1; + +-- Window functions are not allowed in the WHERE clause (pipe operators or otherwise). +table t +|> where first_value(x) over (partition by y) = 1; + +select * from t where first_value(x) over (partition by y) = 1; + +-- Pipe operators may only refer to attributes produced as output from the directly-preceding +-- pipe operator, not from earlier ones. +table t +|> select x, length(y) as z +|> where x + length(y) < 4; + +-- If the WHERE clause wants to filter rows produced by an aggregation, it is not valid to try to +-- refer to the aggregate functions directly; it is necessary to use aliases instead. +(select x, sum(length(y)) as sum_len from t group by x) +|> where sum(length(y)) = 3; + -- Cleanup. ----------- drop table t; diff --git a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out index 7e0b7912105c2..38436b0941034 100644 --- a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out @@ -238,6 +238,56 @@ struct 1 def +-- !query +table t +|> select * +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> select * except (y) +-- !query schema +struct +-- !query output +0 +1 + + +-- !query +table t +|> select /*+ repartition(3) */ * +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> select /*+ repartition(3) */ distinct x +-- !query schema +struct +-- !query output +0 +1 + + +-- !query +table t +|> select /*+ repartition(3) */ all x +-- !query schema +struct +-- !query output +0 +1 + + -- !query table t |> select sum(x) as result @@ -284,6 +334,224 @@ org.apache.spark.sql.AnalysisException } +-- !query +table t +|> where true +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> where x + length(y) < 4 +-- !query schema +struct +-- !query output +0 abc + + +-- !query +table t +|> where x + length(y) < 4 +|> where x + length(y) < 3 +-- !query schema +struct +-- !query output + + + +-- !query +(select x, sum(length(y)) as sum_len from t group by x) +|> where x = 1 +-- !query schema +struct +-- !query output +1 3 + + +-- !query +table t +|> where t.x = 1 +-- !query schema +struct +-- !query output +1 def + + +-- !query +table t +|> where spark_catalog.default.t.x = 1 +-- !query schema +struct +-- !query output +1 def + + +-- !query +(select col from st) +|> where col.i1 = 1 +-- !query schema +struct> +-- !query output + + + +-- !query +table st +|> where st.col.i1 = 2 +-- !query schema +struct> +-- !query output +1 {"i1":2,"i2":3} + + +-- !query +table t +|> where exists (select a from other where x = a limit 1) +-- !query schema +struct +-- !query output +1 def + + +-- !query +table t +|> where (select any_value(a) from other where x = a limit 1) = 1 +-- !query schema +struct +-- !query output +1 def + + +-- !query +table t +|> where sum(x) = 1 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_WHERE_CONDITION", + "sqlState" : "42903", + "messageParameters" : { + "condition" : "\"(sum(x) = 1)\"", + "expressionList" : "sum(spark_catalog.default.t.x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 27, + "fragment" : "table t\n|> where sum(x) = 1" + } ] +} + + +-- !query +table t +|> where y = 'abc' or length(y) + sum(x) = 1 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_WHERE_CONDITION", + "sqlState" : "42903", + "messageParameters" : { + "condition" : "\"((y = abc) OR ((length(y) + sum(x)) = 1))\"", + "expressionList" : "sum(spark_catalog.default.t.x)" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 52, + "fragment" : "table t\n|> where y = 'abc' or length(y) + sum(x) = 1" + } ] +} + + +-- !query +table t +|> where first_value(x) over (partition by y) = 1 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_1034", + "messageParameters" : { + "clauseName" : "WHERE" + } +} + + +-- !query +select * from t where first_value(x) over (partition by y) = 1 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_1034", + "messageParameters" : { + "clauseName" : "WHERE" + } +} + + +-- !query +table t +|> select x, length(y) as z +|> where x + length(y) < 4 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`y`", + "proposal" : "`x`, `z`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 57, + "stopIndex" : 57, + "fragment" : "y" + } ] +} + + +-- !query +(select x, sum(length(y)) as sum_len from t group by x) +|> where sum(length(y)) = 3 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`y`", + "proposal" : "`x`, `sum_len`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 77, + "stopIndex" : 77, + "fragment" : "y" + } ] +} + + -- !query drop table t -- !query schema diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala index a80444feb68ae..ab949c5a21e44 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.analysis.{AnalysisTest, UnresolvedAlias, Un import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Concat, GreaterThan, Literal, NullsFirst, SortOrder, UnresolvedWindowExpression, UnspecifiedFrame, WindowSpecDefinition, WindowSpecReference} import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.trees.TreePattern.{LOCAL_RELATION, PROJECT, UNRESOLVED_RELATION} +import org.apache.spark.sql.catalyst.trees.TreePattern.{FILTER, LOCAL_RELATION, PROJECT, UNRESOLVED_RELATION} import org.apache.spark.sql.connector.catalog.TableCatalog import org.apache.spark.sql.execution.command._ import org.apache.spark.sql.execution.datasources.{CreateTempViewUsing, RefreshResource} @@ -895,6 +895,16 @@ class SparkSqlParserSuite extends AnalysisTest with SharedSparkSession { checkPipeSelect("TABLE t |> SELECT 1 AS X") checkPipeSelect("TABLE t |> SELECT 1 AS X, 2 AS Y |> SELECT X + Y AS Z") checkPipeSelect("VALUES (0), (1) tab(col) |> SELECT col * 2 AS result") + // Basic WHERE operators. + def checkPipeWhere(query: String): Unit = { + val plan: LogicalPlan = parser.parsePlan(query) + assert(plan.containsPattern(FILTER)) + assert(plan.containsAnyPattern(UNRESOLVED_RELATION, LOCAL_RELATION)) + } + checkPipeWhere("TABLE t |> WHERE X = 1") + checkPipeWhere("TABLE t |> SELECT X, LENGTH(Y) AS Z |> WHERE X + LENGTH(Y) < 4") + checkPipeWhere("TABLE t |> WHERE X = 1 AND Y = 2 |> WHERE X + Y = 3") + checkPipeWhere("VALUES (0), (1) tab(col) |> WHERE col < 1") } } } From 70bd606cc865c3d27808eacad85fcf878c23e3a1 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Sat, 21 Sep 2024 11:50:30 +0900 Subject: [PATCH 093/189] [SPARK-49641][DOCS] Include `table_funcs` and `variant_funcs` in the built-in function list doc ### What changes were proposed in this pull request? The pr aims to include `table_funcs` and `variant_funcs` in the built-in function list doc. ### Why are the changes needed? I found that some functions were not involved in our docs, such as `sql_keywords()`, `variant_explode`, etc. Let's include them to improve the user experience for end-users. ### Does this PR introduce _any_ user-facing change? Yes, only for sql api docs. ### How was this patch tested? - Pass GA - Manually check. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48106 from panbingkun/SPARK-49641. Authored-by: panbingkun Signed-off-by: Hyukjin Kwon --- docs/sql-ref-functions-builtin.md | 10 ++++++++++ .../plans/logical/basicLogicalOperators.scala | 15 +++++++++++---- .../spark/sql/api/python/PythonSQLUtils.scala | 7 +++++-- .../org/apache/spark/sql/SQLQuerySuite.scala | 5 ++++- sql/gen-sql-functions-docs.py | 1 + 5 files changed, 31 insertions(+), 7 deletions(-) diff --git a/docs/sql-ref-functions-builtin.md b/docs/sql-ref-functions-builtin.md index c5f4e44dec0d9..b6572609a34b8 100644 --- a/docs/sql-ref-functions-builtin.md +++ b/docs/sql-ref-functions-builtin.md @@ -116,3 +116,13 @@ license: | {% include_api_gen generated-generator-funcs-table.html %} #### Examples {% include_api_gen generated-generator-funcs-examples.html %} + +### Table Functions +{% include_api_gen generated-table-funcs-table.html %} +#### Examples +{% include_api_gen generated-table-funcs-examples.html %} + +### Variant Functions +{% include_api_gen generated-variant-funcs-table.html %} +#### Examples +{% include_api_gen generated-variant-funcs-examples.html %} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala index 926027df4c74b..90af6333b2e0b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala @@ -992,12 +992,18 @@ object Range { castAndEval[Int](expression, IntegerType, paramIndex, paramName) } +// scalastyle:off line.size.limit @ExpressionDescription( usage = """ - _FUNC_(start: long, end: long, step: long, numSlices: integer) - _FUNC_(start: long, end: long, step: long) - _FUNC_(start: long, end: long) - _FUNC_(end: long)""", + _FUNC_(start[, end[, step[, numSlices]]]) / _FUNC_(end) - Returns a table of values within a specified range. + """, + arguments = """ + Arguments: + * start - An optional BIGINT literal defaulted to 0, marking the first value generated. + * end - A BIGINT literal marking endpoint (exclusive) of the number generation. + * step - An optional BIGINT literal defaulted to 1, specifying the increment used when generating values. + * numParts - An optional INTEGER literal specifying how the production of rows is spread across partitions. + """, examples = """ Examples: > SELECT * FROM _FUNC_(1); @@ -1023,6 +1029,7 @@ object Range { """, since = "2.0.0", group = "table_funcs") +// scalastyle:on line.size.limit case class Range( start: Long, end: Long, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala index 93082740cca64..bc270e6ac64ad 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala @@ -29,7 +29,7 @@ import org.apache.spark.internal.LogKeys.CLASS_LOADER import org.apache.spark.security.SocketAuthServer import org.apache.spark.sql.{internal, Column, DataFrame, Row, SparkSession} import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} -import org.apache.spark.sql.catalyst.analysis.FunctionRegistry +import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, TableFunctionRegistry} import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser @@ -69,7 +69,10 @@ private[sql] object PythonSQLUtils extends Logging { // This is needed when generating SQL documentation for built-in functions. def listBuiltinFunctionInfos(): Array[ExpressionInfo] = { - FunctionRegistry.functionSet.flatMap(f => FunctionRegistry.builtin.lookupFunction(f)).toArray + (FunctionRegistry.functionSet.flatMap(f => FunctionRegistry.builtin.lookupFunction(f)) ++ + TableFunctionRegistry.functionSet.flatMap( + f => TableFunctionRegistry.builtin.lookupFunction(f))). + groupBy(_.getName).map(v => v._2.head).toArray } private def listAllSQLConfigs(): Seq[(String, String, String, String)] = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index ce88f7dc475d6..8176d02dbd02d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -111,10 +111,13 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark } test("SPARK-34678: describe functions for table-valued functions") { + sql("describe function range").show(false) checkKeywordsExist(sql("describe function range"), "Function: range", "Class: org.apache.spark.sql.catalyst.plans.logical.Range", - "range(end: long)" + "range(start[, end[, step[, numSlices]]])", + "range(end)", + "Returns a table of values within a specified range." ) } diff --git a/sql/gen-sql-functions-docs.py b/sql/gen-sql-functions-docs.py index 4be9966747d1f..a1facbaaf7e3b 100644 --- a/sql/gen-sql-functions-docs.py +++ b/sql/gen-sql-functions-docs.py @@ -36,6 +36,7 @@ "bitwise_funcs", "conversion_funcs", "csv_funcs", "xml_funcs", "lambda_funcs", "collection_funcs", "url_funcs", "hash_funcs", "struct_funcs", + "table_funcs", "variant_funcs" } From f235bab24761d8049e3d74411c19ddf3e3b5a697 Mon Sep 17 00:00:00 2001 From: Harsh Motwani Date: Sat, 21 Sep 2024 12:27:05 +0900 Subject: [PATCH 094/189] [SPARK-49451][FOLLOW-UP] Add support for duplicate keys in from_json(_, 'variant') ### What changes were proposed in this pull request? This PR adds support for duplicate key support in the `from_json(_, 'variant')` query pattern. Duplicate key support [has been introduced](https://github.com/apache/spark/pull/47920) in `parse_json`, json scans and the `from_json` expressions with nested schemas but this code path was not updated. ### Why are the changes needed? This change makes the behavior of `from_json(_, 'variant')` consistent with every other variant construction expression. ### Does this PR introduce _any_ user-facing change? It potentially allows users to use the `from_json(, 'variant')` expression on json inputs with duplicate keys depending on a config. ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48177 from harshmotw-db/harshmotw-db/master. Authored-by: Harsh Motwani Signed-off-by: Hyukjin Kwon --- .../expressions/jsonExpressions.scala | 12 +++++-- .../function_from_json.explain | 2 +- .../function_from_json_orphaned.explain | 2 +- ...unction_from_json_with_json_schema.explain | 2 +- .../analyzer-results/ansi/date.sql.out | 2 +- .../analyzer-results/ansi/interval.sql.out | 6 ++-- .../ansi/parse-schema-string.sql.out | 4 +-- .../analyzer-results/ansi/timestamp.sql.out | 2 +- .../sql-tests/analyzer-results/date.sql.out | 2 +- .../analyzer-results/datetime-legacy.sql.out | 4 +-- .../analyzer-results/interval.sql.out | 6 ++-- .../analyzer-results/json-functions.sql.out | 34 +++++++++---------- .../parse-schema-string.sql.out | 4 +-- .../sql-session-variables.sql.out | 2 +- .../subexp-elimination.sql.out | 10 +++--- .../analyzer-results/timestamp.sql.out | 2 +- .../timestampNTZ/timestamp-ansi.sql.out | 2 +- .../timestampNTZ/timestamp.sql.out | 2 +- .../native/stringCastAndExpressions.sql.out | 2 +- .../spark/sql/VariantEndToEndSuite.scala | 32 +++++++++++++++++ 20 files changed, 87 insertions(+), 47 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala index 574a61cf9c903..2037eb22fede6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala @@ -632,7 +632,8 @@ case class JsonToStructs( schema: DataType, options: Map[String, String], child: Expression, - timeZoneId: Option[String] = None) + timeZoneId: Option[String] = None, + variantAllowDuplicateKeys: Boolean = SQLConf.get.getConf(SQLConf.VARIANT_ALLOW_DUPLICATE_KEYS)) extends UnaryExpression with TimeZoneAwareExpression with CodegenFallback @@ -719,7 +720,8 @@ case class JsonToStructs( override def nullSafeEval(json: Any): Any = nullableSchema match { case _: VariantType => - VariantExpressionEvalUtils.parseJson(json.asInstanceOf[UTF8String]) + VariantExpressionEvalUtils.parseJson(json.asInstanceOf[UTF8String], + allowDuplicateKeys = variantAllowDuplicateKeys) case _ => converter(parser.parse(json.asInstanceOf[UTF8String])) } @@ -737,6 +739,12 @@ case class JsonToStructs( copy(child = newChild) } +object JsonToStructs { + def unapply( + j: JsonToStructs): Option[(DataType, Map[String, String], Expression, Option[String])] = + Some((j.schema, j.options, j.child, j.timeZoneId)) +} + /** * Converts a [[StructType]], [[ArrayType]] or [[MapType]] to a JSON output string. */ diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json.explain index 1219f11d4696e..8d1d122d156ff 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json.explain @@ -1,2 +1,2 @@ -Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles)) AS from_json(g)#0] +Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles), false) AS from_json(g)#0] +- LocalRelation , [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_orphaned.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_orphaned.explain index 1219f11d4696e..8d1d122d156ff 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_orphaned.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_orphaned.explain @@ -1,2 +1,2 @@ -Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles)) AS from_json(g)#0] +Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles), false) AS from_json(g)#0] +- LocalRelation , [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_with_json_schema.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_with_json_schema.explain index 1219f11d4696e..8d1d122d156ff 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_with_json_schema.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_json_with_json_schema.explain @@ -1,2 +1,2 @@ -Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles)) AS from_json(g)#0] +Project [from_json(StructField(id,LongType,true), StructField(a,IntegerType,true), StructField(b,DoubleType,true), g#0, Some(America/Los_Angeles), false) AS from_json(g)#0] +- LocalRelation , [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/date.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/date.sql.out index fd927b99c6456..0e4d2d4e99e26 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/date.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/date.sql.out @@ -736,7 +736,7 @@ Project [to_date(26/October/2015, Some(dd/MMMMM/yyyy), Some(America/Los_Angeles) -- !query select from_json('{"d":"26/October/2015"}', 'd Date', map('dateFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"d":"26/October/2015"})#x] +Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"d":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/interval.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/interval.sql.out index 472c9b1df064a..b0d128c4cab69 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/interval.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/interval.sql.out @@ -2108,7 +2108,7 @@ SELECT to_csv(named_struct('a', interval 32 year, 'b', interval 10 month)), from_csv(to_csv(named_struct('a', interval 32 year, 'b', interval 10 month)), 'a interval year, b interval month') -- !query analysis -Project [from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles)) AS from_json({"a":"1 days"})#x, from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None) AS from_csv(1, 1)#x, to_json(from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1 days"}))#x, to_csv(from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None), Some(America/Los_Angeles)) AS to_csv(from_csv(1, 1))#x, to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)) AS to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH))#x, from_csv(StructField(a,YearMonthIntervalType(0,0),true), StructField(b,YearMonthIntervalType(1,1),true), to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), None) AS from_csv(to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH)))#x] +Project [from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles), false) AS from_json({"a":"1 days"})#x, from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None) AS from_csv(1, 1)#x, to_json(from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1 days"}))#x, to_csv(from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None), Some(America/Los_Angeles)) AS to_csv(from_csv(1, 1))#x, to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)) AS to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH))#x, from_csv(StructField(a,YearMonthIntervalType(0,0),true), StructField(b,YearMonthIntervalType(1,1),true), to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), None) AS from_csv(to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH)))#x] +- OneRowRelation @@ -2119,7 +2119,7 @@ SELECT to_json(map('a', interval 100 day 130 minute)), from_json(to_json(map('a', interval 100 day 130 minute)), 'a interval day to minute') -- !query analysis -Project [from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE))#x, from_json(StructField(a,DayTimeIntervalType(0,2),true), to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS from_json(to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE)))#x] +Project [from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE))#x, from_json(StructField(a,DayTimeIntervalType(0,2),true), to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)), Some(America/Los_Angeles), false) AS from_json(to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE)))#x] +- OneRowRelation @@ -2130,7 +2130,7 @@ SELECT to_json(map('a', interval 32 year 10 month)), from_json(to_json(map('a', interval 32 year 10 month)), 'a interval year to month') -- !query analysis -Project [from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '32-10' YEAR TO MONTH))#x, from_json(StructField(a,YearMonthIntervalType(0,1),true), to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS from_json(to_json(map(a, INTERVAL '32-10' YEAR TO MONTH)))#x] +Project [from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '32-10' YEAR TO MONTH))#x, from_json(StructField(a,YearMonthIntervalType(0,1),true), to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), false) AS from_json(to_json(map(a, INTERVAL '32-10' YEAR TO MONTH)))#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/parse-schema-string.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/parse-schema-string.sql.out index 45fc3bd03a782..ae8e47ed3665c 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/parse-schema-string.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/parse-schema-string.sql.out @@ -16,12 +16,12 @@ Project [from_csv(StructField(cube,IntegerType,true), 1, Some(America/Los_Angele -- !query select from_json('{"create":1}', 'create INT') -- !query analysis -Project [from_json(StructField(create,IntegerType,true), {"create":1}, Some(America/Los_Angeles)) AS from_json({"create":1})#x] +Project [from_json(StructField(create,IntegerType,true), {"create":1}, Some(America/Los_Angeles), false) AS from_json({"create":1})#x] +- OneRowRelation -- !query select from_json('{"cube":1}', 'cube INT') -- !query analysis -Project [from_json(StructField(cube,IntegerType,true), {"cube":1}, Some(America/Los_Angeles)) AS from_json({"cube":1})#x] +Project [from_json(StructField(cube,IntegerType,true), {"cube":1}, Some(America/Los_Angeles), false) AS from_json({"cube":1})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/timestamp.sql.out index bf34490d657e3..560974d28c545 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/ansi/timestamp.sql.out @@ -730,7 +730,7 @@ Project [unix_timestamp(22 05 2020 Friday, dd MM yyyy EEEEE, Some(America/Los_An -- !query select from_json('{"t":"26/October/2015"}', 't Timestamp', map('timestampFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"t":"26/October/2015"})#x] +Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"t":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/date.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/date.sql.out index 48137e06467e8..88c7d7b4e7d72 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/date.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/date.sql.out @@ -811,7 +811,7 @@ Project [to_date(26/October/2015, Some(dd/MMMMM/yyyy), Some(America/Los_Angeles) -- !query select from_json('{"d":"26/October/2015"}', 'd Date', map('dateFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"d":"26/October/2015"})#x] +Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"d":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/datetime-legacy.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/datetime-legacy.sql.out index 1e49f4df8267a..4221db822d024 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/datetime-legacy.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/datetime-legacy.sql.out @@ -811,7 +811,7 @@ Project [to_date(26/October/2015, Some(dd/MMMMM/yyyy), Some(America/Los_Angeles) -- !query select from_json('{"d":"26/October/2015"}', 'd Date', map('dateFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"d":"26/October/2015"})#x] +Project [from_json(StructField(d,DateType,true), (dateFormat,dd/MMMMM/yyyy), {"d":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"d":"26/October/2015"})#x] +- OneRowRelation @@ -1833,7 +1833,7 @@ Project [unix_timestamp(22 05 2020 Friday, dd MM yyyy EEEEE, Some(America/Los_An -- !query select from_json('{"t":"26/October/2015"}', 't Timestamp', map('timestampFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"t":"26/October/2015"})#x] +Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"t":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/interval.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/interval.sql.out index 3db38d482b26d..efa149509751d 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/interval.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/interval.sql.out @@ -2108,7 +2108,7 @@ SELECT to_csv(named_struct('a', interval 32 year, 'b', interval 10 month)), from_csv(to_csv(named_struct('a', interval 32 year, 'b', interval 10 month)), 'a interval year, b interval month') -- !query analysis -Project [from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles)) AS from_json({"a":"1 days"})#x, from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None) AS from_csv(1, 1)#x, to_json(from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1 days"}))#x, to_csv(from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None), Some(America/Los_Angeles)) AS to_csv(from_csv(1, 1))#x, to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)) AS to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH))#x, from_csv(StructField(a,YearMonthIntervalType(0,0),true), StructField(b,YearMonthIntervalType(1,1),true), to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), None) AS from_csv(to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH)))#x] +Project [from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles), false) AS from_json({"a":"1 days"})#x, from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None) AS from_csv(1, 1)#x, to_json(from_json(StructField(a,CalendarIntervalType,true), {"a":"1 days"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1 days"}))#x, to_csv(from_csv(StructField(a,IntegerType,true), StructField(b,YearMonthIntervalType(0,0),true), 1, 1, Some(America/Los_Angeles), None), Some(America/Los_Angeles)) AS to_csv(from_csv(1, 1))#x, to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)) AS to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH))#x, from_csv(StructField(a,YearMonthIntervalType(0,0),true), StructField(b,YearMonthIntervalType(1,1),true), to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), None) AS from_csv(to_csv(named_struct(a, INTERVAL '32' YEAR, b, INTERVAL '10' MONTH)))#x] +- OneRowRelation @@ -2119,7 +2119,7 @@ SELECT to_json(map('a', interval 100 day 130 minute)), from_json(to_json(map('a', interval 100 day 130 minute)), 'a interval day to minute') -- !query analysis -Project [from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE))#x, from_json(StructField(a,DayTimeIntervalType(0,2),true), to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS from_json(to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE)))#x] +Project [from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,DayTimeIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE))#x, from_json(StructField(a,DayTimeIntervalType(0,2),true), to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE), Some(America/Los_Angeles)), Some(America/Los_Angeles), false) AS from_json(to_json(map(a, INTERVAL '100 02:10' DAY TO MINUTE)))#x] +- OneRowRelation @@ -2130,7 +2130,7 @@ SELECT to_json(map('a', interval 32 year 10 month)), from_json(to_json(map('a', interval 32 year 10 month)), 'a interval year to month') -- !query analysis -Project [from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '32-10' YEAR TO MONTH))#x, from_json(StructField(a,YearMonthIntervalType(0,1),true), to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles)) AS from_json(to_json(map(a, INTERVAL '32-10' YEAR TO MONTH)))#x] +Project [from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false) AS from_json({"a":"1"})#x, to_json(from_json(StructField(a,YearMonthIntervalType(0,0),true), {"a":"1"}, Some(America/Los_Angeles), false), Some(America/Los_Angeles)) AS to_json(from_json({"a":"1"}))#x, to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)) AS to_json(map(a, INTERVAL '32-10' YEAR TO MONTH))#x, from_json(StructField(a,YearMonthIntervalType(0,1),true), to_json(map(a, INTERVAL '32-10' YEAR TO MONTH), Some(America/Los_Angeles)), Some(America/Los_Angeles), false) AS from_json(to_json(map(a, INTERVAL '32-10' YEAR TO MONTH)))#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out index 0d7c6b2056231..fef9d0c5b6250 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out @@ -118,14 +118,14 @@ org.apache.spark.sql.AnalysisException -- !query select from_json('{"a":1}', 'a INT') -- !query analysis -Project [from_json(StructField(a,IntegerType,true), {"a":1}, Some(America/Los_Angeles)) AS from_json({"a":1})#x] +Project [from_json(StructField(a,IntegerType,true), {"a":1}, Some(America/Los_Angeles), false) AS from_json({"a":1})#x] +- OneRowRelation -- !query select from_json('{"time":"26/08/2015"}', 'time Timestamp', map('timestampFormat', 'dd/MM/yyyy')) -- !query analysis -Project [from_json(StructField(time,TimestampType,true), (timestampFormat,dd/MM/yyyy), {"time":"26/08/2015"}, Some(America/Los_Angeles)) AS from_json({"time":"26/08/2015"})#x] +Project [from_json(StructField(time,TimestampType,true), (timestampFormat,dd/MM/yyyy), {"time":"26/08/2015"}, Some(America/Los_Angeles), false) AS from_json({"time":"26/08/2015"})#x] +- OneRowRelation @@ -279,14 +279,14 @@ DropTempViewCommand jsonTable -- !query select from_json('{"a":1, "b":2}', 'map') -- !query analysis -Project [from_json(MapType(StringType,IntegerType,true), {"a":1, "b":2}, Some(America/Los_Angeles)) AS entries#x] +Project [from_json(MapType(StringType,IntegerType,true), {"a":1, "b":2}, Some(America/Los_Angeles), false) AS entries#x] +- OneRowRelation -- !query select from_json('{"a":1, "b":"2"}', 'struct') -- !query analysis -Project [from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), {"a":1, "b":"2"}, Some(America/Los_Angeles)) AS from_json({"a":1, "b":"2"})#x] +Project [from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), {"a":1, "b":"2"}, Some(America/Los_Angeles), false) AS from_json({"a":1, "b":"2"})#x] +- OneRowRelation @@ -300,70 +300,70 @@ Project [schema_of_json({"c1":0, "c2":[1]}) AS schema_of_json({"c1":0, "c2":[1]} -- !query select from_json('{"c1":[1, 2, 3]}', schema_of_json('{"c1":[0]}')) -- !query analysis -Project [from_json(StructField(c1,ArrayType(LongType,true),true), {"c1":[1, 2, 3]}, Some(America/Los_Angeles)) AS from_json({"c1":[1, 2, 3]})#x] +Project [from_json(StructField(c1,ArrayType(LongType,true),true), {"c1":[1, 2, 3]}, Some(America/Los_Angeles), false) AS from_json({"c1":[1, 2, 3]})#x] +- OneRowRelation -- !query select from_json('[1, 2, 3]', 'array') -- !query analysis -Project [from_json(ArrayType(IntegerType,true), [1, 2, 3], Some(America/Los_Angeles)) AS from_json([1, 2, 3])#x] +Project [from_json(ArrayType(IntegerType,true), [1, 2, 3], Some(America/Los_Angeles), false) AS from_json([1, 2, 3])#x] +- OneRowRelation -- !query select from_json('[1, "2", 3]', 'array') -- !query analysis -Project [from_json(ArrayType(IntegerType,true), [1, "2", 3], Some(America/Los_Angeles)) AS from_json([1, "2", 3])#x] +Project [from_json(ArrayType(IntegerType,true), [1, "2", 3], Some(America/Los_Angeles), false) AS from_json([1, "2", 3])#x] +- OneRowRelation -- !query select from_json('[1, 2, null]', 'array') -- !query analysis -Project [from_json(ArrayType(IntegerType,true), [1, 2, null], Some(America/Los_Angeles)) AS from_json([1, 2, null])#x] +Project [from_json(ArrayType(IntegerType,true), [1, 2, null], Some(America/Los_Angeles), false) AS from_json([1, 2, null])#x] +- OneRowRelation -- !query select from_json('[{"a": 1}, {"a":2}]', 'array>') -- !query analysis -Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), [{"a": 1}, {"a":2}], Some(America/Los_Angeles)) AS from_json([{"a": 1}, {"a":2}])#x] +Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), [{"a": 1}, {"a":2}], Some(America/Los_Angeles), false) AS from_json([{"a": 1}, {"a":2}])#x] +- OneRowRelation -- !query select from_json('{"a": 1}', 'array>') -- !query analysis -Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), {"a": 1}, Some(America/Los_Angeles)) AS from_json({"a": 1})#x] +Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), {"a": 1}, Some(America/Los_Angeles), false) AS from_json({"a": 1})#x] +- OneRowRelation -- !query select from_json('[null, {"a":2}]', 'array>') -- !query analysis -Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), [null, {"a":2}], Some(America/Los_Angeles)) AS from_json([null, {"a":2}])#x] +Project [from_json(ArrayType(StructType(StructField(a,IntegerType,true)),true), [null, {"a":2}], Some(America/Los_Angeles), false) AS from_json([null, {"a":2}])#x] +- OneRowRelation -- !query select from_json('[{"a": 1}, {"b":2}]', 'array>') -- !query analysis -Project [from_json(ArrayType(MapType(StringType,IntegerType,true),true), [{"a": 1}, {"b":2}], Some(America/Los_Angeles)) AS from_json([{"a": 1}, {"b":2}])#x] +Project [from_json(ArrayType(MapType(StringType,IntegerType,true),true), [{"a": 1}, {"b":2}], Some(America/Los_Angeles), false) AS from_json([{"a": 1}, {"b":2}])#x] +- OneRowRelation -- !query select from_json('[{"a": 1}, 2]', 'array>') -- !query analysis -Project [from_json(ArrayType(MapType(StringType,IntegerType,true),true), [{"a": 1}, 2], Some(America/Los_Angeles)) AS from_json([{"a": 1}, 2])#x] +Project [from_json(ArrayType(MapType(StringType,IntegerType,true),true), [{"a": 1}, 2], Some(America/Los_Angeles), false) AS from_json([{"a": 1}, 2])#x] +- OneRowRelation -- !query select from_json('{"d": "2012-12-15", "t": "2012-12-15 15:15:15"}', 'd date, t timestamp') -- !query analysis -Project [from_json(StructField(d,DateType,true), StructField(t,TimestampType,true), {"d": "2012-12-15", "t": "2012-12-15 15:15:15"}, Some(America/Los_Angeles)) AS from_json({"d": "2012-12-15", "t": "2012-12-15 15:15:15"})#x] +Project [from_json(StructField(d,DateType,true), StructField(t,TimestampType,true), {"d": "2012-12-15", "t": "2012-12-15 15:15:15"}, Some(America/Los_Angeles), false) AS from_json({"d": "2012-12-15", "t": "2012-12-15 15:15:15"})#x] +- OneRowRelation @@ -373,7 +373,7 @@ select from_json( 'd date, t timestamp', map('dateFormat', 'MM/dd yyyy', 'timestampFormat', 'MM/dd yyyy HH:mm:ss')) -- !query analysis -Project [from_json(StructField(d,DateType,true), StructField(t,TimestampType,true), (dateFormat,MM/dd yyyy), (timestampFormat,MM/dd yyyy HH:mm:ss), {"d": "12/15 2012", "t": "12/15 2012 15:15:15"}, Some(America/Los_Angeles)) AS from_json({"d": "12/15 2012", "t": "12/15 2012 15:15:15"})#x] +Project [from_json(StructField(d,DateType,true), StructField(t,TimestampType,true), (dateFormat,MM/dd yyyy), (timestampFormat,MM/dd yyyy HH:mm:ss), {"d": "12/15 2012", "t": "12/15 2012 15:15:15"}, Some(America/Los_Angeles), false) AS from_json({"d": "12/15 2012", "t": "12/15 2012 15:15:15"})#x] +- OneRowRelation @@ -383,7 +383,7 @@ select from_json( 'd date', map('dateFormat', 'MM-dd')) -- !query analysis -Project [from_json(StructField(d,DateType,true), (dateFormat,MM-dd), {"d": "02-29"}, Some(America/Los_Angeles)) AS from_json({"d": "02-29"})#x] +Project [from_json(StructField(d,DateType,true), (dateFormat,MM-dd), {"d": "02-29"}, Some(America/Los_Angeles), false) AS from_json({"d": "02-29"})#x] +- OneRowRelation @@ -393,7 +393,7 @@ select from_json( 't timestamp', map('timestampFormat', 'MM-dd')) -- !query analysis -Project [from_json(StructField(t,TimestampType,true), (timestampFormat,MM-dd), {"t": "02-29"}, Some(America/Los_Angeles)) AS from_json({"t": "02-29"})#x] +Project [from_json(StructField(t,TimestampType,true), (timestampFormat,MM-dd), {"t": "02-29"}, Some(America/Los_Angeles), false) AS from_json({"t": "02-29"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/parse-schema-string.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-schema-string.sql.out index 45fc3bd03a782..ae8e47ed3665c 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/parse-schema-string.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-schema-string.sql.out @@ -16,12 +16,12 @@ Project [from_csv(StructField(cube,IntegerType,true), 1, Some(America/Los_Angele -- !query select from_json('{"create":1}', 'create INT') -- !query analysis -Project [from_json(StructField(create,IntegerType,true), {"create":1}, Some(America/Los_Angeles)) AS from_json({"create":1})#x] +Project [from_json(StructField(create,IntegerType,true), {"create":1}, Some(America/Los_Angeles), false) AS from_json({"create":1})#x] +- OneRowRelation -- !query select from_json('{"cube":1}', 'cube INT') -- !query analysis -Project [from_json(StructField(cube,IntegerType,true), {"cube":1}, Some(America/Los_Angeles)) AS from_json({"cube":1})#x] +Project [from_json(StructField(cube,IntegerType,true), {"cube":1}, Some(America/Los_Angeles), false) AS from_json({"cube":1})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/sql-session-variables.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/sql-session-variables.sql.out index a4e40f08b4463..02e7c39ae83fd 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/sql-session-variables.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/sql-session-variables.sql.out @@ -2147,7 +2147,7 @@ CreateVariable defaultvalueexpression(cast(a INT as string), 'a INT'), true -- !query SELECT from_json('{"a": 1}', var1) -- !query analysis -Project [from_json(StructField(a,IntegerType,true), {"a": 1}, Some(America/Los_Angeles)) AS from_json({"a": 1})#x] +Project [from_json(StructField(a,IntegerType,true), {"a": 1}, Some(America/Los_Angeles), false) AS from_json({"a": 1})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/subexp-elimination.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/subexp-elimination.sql.out index 94073f2751b3e..754b05bfa6fed 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/subexp-elimination.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/subexp-elimination.sql.out @@ -15,7 +15,7 @@ AS testData(a, b), false, true, LocalTempView, UNSUPPORTED, true -- !query SELECT from_json(a, 'struct').a, from_json(a, 'struct').b, from_json(b, 'array>')[0].a, from_json(b, 'array>')[0].b FROM testData -- !query analysis -Project [from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a AS from_json(a).a#x, from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).b AS from_json(a).b#x, from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].a AS from_json(b)[0].a#x, from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b AS from_json(b)[0].b#x] +Project [from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a AS from_json(a).a#x, from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).b AS from_json(a).b#x, from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].a AS from_json(b)[0].a#x, from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b AS from_json(b)[0].b#x] +- SubqueryAlias testdata +- View (`testData`, [a#x, b#x]) +- Project [cast(a#x as string) AS a#x, cast(b#x as string) AS b#x] @@ -27,7 +27,7 @@ Project [from_json(StructField(a,IntegerType,true), StructField(b,StringType,tru -- !query SELECT if(from_json(a, 'struct').a > 1, from_json(b, 'array>')[0].a, from_json(b, 'array>')[0].a + 1) FROM testData -- !query analysis -Project [if ((from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a > 1)) from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].a else (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].a + 1) AS (IF((from_json(a).a > 1), from_json(b)[0].a, (from_json(b)[0].a + 1)))#x] +Project [if ((from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a > 1)) from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].a else (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].a + 1) AS (IF((from_json(a).a > 1), from_json(b)[0].a, (from_json(b)[0].a + 1)))#x] +- SubqueryAlias testdata +- View (`testData`, [a#x, b#x]) +- Project [cast(a#x as string) AS a#x, cast(b#x as string) AS b#x] @@ -39,7 +39,7 @@ Project [if ((from_json(StructField(a,IntegerType,true), StructField(b,StringTyp -- !query SELECT if(isnull(from_json(a, 'struct').a), from_json(b, 'array>')[0].b + 1, from_json(b, 'array>')[0].b) FROM testData -- !query analysis -Project [if (isnull(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a)) (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b + 1) else from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b AS (IF((from_json(a).a IS NULL), (from_json(b)[0].b + 1), from_json(b)[0].b))#x] +Project [if (isnull(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a)) (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b + 1) else from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b AS (IF((from_json(a).a IS NULL), (from_json(b)[0].b + 1), from_json(b)[0].b))#x] +- SubqueryAlias testdata +- View (`testData`, [a#x, b#x]) +- Project [cast(a#x as string) AS a#x, cast(b#x as string) AS b#x] @@ -51,7 +51,7 @@ Project [if (isnull(from_json(StructField(a,IntegerType,true), StructField(b,Str -- !query SELECT case when from_json(a, 'struct').a > 5 then from_json(a, 'struct').b when from_json(a, 'struct').a > 4 then from_json(a, 'struct').b + 1 else from_json(a, 'struct').b + 2 end FROM testData -- !query analysis -Project [CASE WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a > 5) THEN from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).b WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a > 4) THEN cast((cast(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).b as double) + cast(1 as double)) as string) ELSE cast((cast(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).b as double) + cast(2 as double)) as string) END AS CASE WHEN (from_json(a).a > 5) THEN from_json(a).b WHEN (from_json(a).a > 4) THEN (from_json(a).b + 1) ELSE (from_json(a).b + 2) END#x] +Project [CASE WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a > 5) THEN from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).b WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a > 4) THEN cast((cast(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).b as double) + cast(1 as double)) as string) ELSE cast((cast(from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).b as double) + cast(2 as double)) as string) END AS CASE WHEN (from_json(a).a > 5) THEN from_json(a).b WHEN (from_json(a).a > 4) THEN (from_json(a).b + 1) ELSE (from_json(a).b + 2) END#x] +- SubqueryAlias testdata +- View (`testData`, [a#x, b#x]) +- Project [cast(a#x as string) AS a#x, cast(b#x as string) AS b#x] @@ -63,7 +63,7 @@ Project [CASE WHEN (from_json(StructField(a,IntegerType,true), StructField(b,Str -- !query SELECT case when from_json(a, 'struct').a > 5 then from_json(b, 'array>')[0].b when from_json(a, 'struct').a > 4 then from_json(b, 'array>')[0].b + 1 else from_json(b, 'array>')[0].b + 2 end FROM testData -- !query analysis -Project [CASE WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a > 5) THEN from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles)).a > 4) THEN (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b + 1) ELSE (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles))[0].b + 2) END AS CASE WHEN (from_json(a).a > 5) THEN from_json(b)[0].b WHEN (from_json(a).a > 4) THEN (from_json(b)[0].b + 1) ELSE (from_json(b)[0].b + 2) END#x] +Project [CASE WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a > 5) THEN from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b WHEN (from_json(StructField(a,IntegerType,true), StructField(b,StringType,true), a#x, Some(America/Los_Angeles), false).a > 4) THEN (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b + 1) ELSE (from_json(ArrayType(StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)),true), b#x, Some(America/Los_Angeles), false)[0].b + 2) END AS CASE WHEN (from_json(a).a > 5) THEN from_json(b)[0].b WHEN (from_json(a).a > 4) THEN (from_json(b)[0].b + 1) ELSE (from_json(b)[0].b + 2) END#x] +- SubqueryAlias testdata +- View (`testData`, [a#x, b#x]) +- Project [cast(a#x as string) AS a#x, cast(b#x as string) AS b#x] diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp.sql.out index 6ca35b8b141dc..dcfd783b648f8 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp.sql.out @@ -802,7 +802,7 @@ Project [unix_timestamp(22 05 2020 Friday, dd MM yyyy EEEEE, Some(America/Los_An -- !query select from_json('{"t":"26/October/2015"}', 't Timestamp', map('timestampFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"t":"26/October/2015"})#x] +Project [from_json(StructField(t,TimestampType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"t":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp-ansi.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp-ansi.sql.out index e50c860270563..ec227afc87fe1 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp-ansi.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp-ansi.sql.out @@ -745,7 +745,7 @@ Project [unix_timestamp(22 05 2020 Friday, dd MM yyyy EEEEE, Some(America/Los_An -- !query select from_json('{"t":"26/October/2015"}', 't Timestamp', map('timestampFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(t,TimestampNTZType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"t":"26/October/2015"})#x] +Project [from_json(StructField(t,TimestampNTZType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"t":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out index 098abfb3852cf..7475f837250d5 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out @@ -805,7 +805,7 @@ Project [unix_timestamp(22 05 2020 Friday, dd MM yyyy EEEEE, Some(America/Los_An -- !query select from_json('{"t":"26/October/2015"}', 't Timestamp', map('timestampFormat', 'dd/MMMMM/yyyy')) -- !query analysis -Project [from_json(StructField(t,TimestampNTZType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles)) AS from_json({"t":"26/October/2015"})#x] +Project [from_json(StructField(t,TimestampNTZType,true), (timestampFormat,dd/MMMMM/yyyy), {"t":"26/October/2015"}, Some(America/Los_Angeles), false) AS from_json({"t":"26/October/2015"})#x] +- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/stringCastAndExpressions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/stringCastAndExpressions.sql.out index 009e91f7ffacf..22e60d0606382 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/stringCastAndExpressions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/stringCastAndExpressions.sql.out @@ -370,7 +370,7 @@ Project [c0#x] -- !query select from_json(a, 'a INT') from t -- !query analysis -Project [from_json(StructField(a,IntegerType,true), a#x, Some(America/Los_Angeles)) AS from_json(a)#x] +Project [from_json(StructField(a,IntegerType,true), a#x, Some(America/Los_Angeles), false) AS from_json(a)#x] +- SubqueryAlias t +- View (`t`, [a#x]) +- Project [cast(a#x as string) AS a#x] diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala index 3224baf42f3e5..19d4ac23709b6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala @@ -16,6 +16,7 @@ */ package org.apache.spark.sql +import org.apache.spark.SparkThrowable import org.apache.spark.sql.QueryTest.sameRows import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} @@ -28,6 +29,7 @@ import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarArray import org.apache.spark.types.variant.VariantBuilder +import org.apache.spark.types.variant.VariantUtil._ import org.apache.spark.unsafe.types.VariantVal class VariantEndToEndSuite extends QueryTest with SharedSparkSession { @@ -37,8 +39,10 @@ class VariantEndToEndSuite extends QueryTest with SharedSparkSession { def check(input: String, output: String = null): Unit = { val df = Seq(input).toDF("v") val variantDF = df.select(to_json(parse_json(col("v")))) + val variantDF2 = df.select(to_json(from_json(col("v"), VariantType))) val expected = if (output != null) output else input checkAnswer(variantDF, Seq(Row(expected))) + checkAnswer(variantDF2, Seq(Row(expected))) } check("null") @@ -339,4 +343,32 @@ class VariantEndToEndSuite extends QueryTest with SharedSparkSession { } } } + + test("from_json(_, 'variant') with duplicate keys") { + val json: String = """{"a": 1, "b": 2, "c": "3", "a": 4}""" + withSQLConf(SQLConf.VARIANT_ALLOW_DUPLICATE_KEYS.key -> "true") { + val df = Seq(json).toDF("j") + .selectExpr("from_json(j,'variant')") + val actual = df.collect().head(0).asInstanceOf[VariantVal] + val expectedValue: Array[Byte] = Array(objectHeader(false, 1, 1), + /* size */ 3, + /* id list */ 0, 1, 2, + /* offset list */ 4, 0, 2, 6, + /* field data */ primitiveHeader(INT1), 2, shortStrHeader(1), '3', + primitiveHeader(INT1), 4) + val expectedMetadata: Array[Byte] = Array(VERSION, 3, 0, 1, 2, 3, 'a', 'b', 'c') + assert(actual === new VariantVal(expectedValue, expectedMetadata)) + } + withSQLConf(SQLConf.VARIANT_ALLOW_DUPLICATE_KEYS.key -> "false") { + val df = Seq(json).toDF("j") + .selectExpr("from_json(j,'variant')") + checkError( + exception = intercept[SparkThrowable] { + df.collect() + }, + condition = "MALFORMED_RECORD_IN_PARSING.WITHOUT_SUGGESTION", + parameters = Map("badRecord" -> json, "failFastMode" -> "FAILFAST") + ) + } + } } From fc8b94544163cd1988053a3f8eb8b4770fbbb55b Mon Sep 17 00:00:00 2001 From: Ziqi Liu Date: Sat, 21 Sep 2024 12:37:57 +0900 Subject: [PATCH 095/189] [SPARK-49460][SQL] Followup: fix potential NPE risk ### What changes were proposed in this pull request? Fixed potential NPE risk in `EmptyRelationExec.logical` ### Why are the changes needed? This is a follow up for https://github.com/apache/spark/pull/47931, I've checked other callsites of `EmptyRelationExec.logical`, which we can not assure it's driver-only. So we should fix those potential risks. ### Does this PR introduce _any_ user-facing change? NO ### How was this patch tested? Existing UT ### Was this patch authored or co-authored using generative AI tooling? NO Closes #48191 from liuzqt/SPARK-49460. Authored-by: Ziqi Liu Signed-off-by: Hyukjin Kwon --- .../spark/sql/execution/EmptyRelationExec.scala | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/EmptyRelationExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/EmptyRelationExec.scala index 8a544de7567e8..a0c3d7b51c2c3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/EmptyRelationExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/EmptyRelationExec.scala @@ -71,13 +71,15 @@ case class EmptyRelationExec(@transient logical: LogicalPlan) extends LeafExecNo maxFields, printNodeId, indent) - lastChildren.add(true) - logical.generateTreeString( - depth + 1, lastChildren, append, verbose, "", false, maxFields, printNodeId, indent) - lastChildren.remove(lastChildren.size() - 1) + Option(logical).foreach { _ => + lastChildren.add(true) + logical.generateTreeString( + depth + 1, lastChildren, append, verbose, "", false, maxFields, printNodeId, indent) + lastChildren.remove(lastChildren.size() - 1) + } } override def doCanonicalize(): SparkPlan = { - this.copy(logical = LocalRelation(logical.output).canonicalized) + this.copy(logical = LocalRelation(output).canonicalized) } } From 0b05b1aa72ced85b49c7230a493bd3200bcc786a Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sat, 21 Sep 2024 09:20:23 +0200 Subject: [PATCH 096/189] [SPARK-48782][SQL][TESTS][FOLLOW-UP] Enable ANSI for malformed input test in ProcedureSuite ### What changes were proposed in this pull request? This PR is a followup of https://github.com/apache/spark/pull/47943 that enables ANSI for malformed input test in ProcedureSuite. ### Why are the changes needed? The specific test fails with ANSI mode disabled https://github.com/apache/spark/actions/runs/10951615244/job/30408963913 ``` - malformed input to implicit cast *** FAILED *** (4 milliseconds) Expected exception org.apache.spark.SparkNumberFormatException to be thrown, but no exception was thrown (ProcedureSuite.scala:264) org.scalatest.exceptions.TestFailedException: at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472) at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471) at org.scalatest.funsuite.AnyFunSuite.newAssertionFailedException(AnyFunSuite.scala:1564) ... ``` The test depends on `sum`'s failure so this PR simply enables ANSI mode for that specific test. ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? Manually ran with ANSI mode off. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48193 from HyukjinKwon/SPARK-48782-followup. Authored-by: Hyukjin Kwon Signed-off-by: Max Gekk --- .../spark/sql/connector/ProcedureSuite.scala | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala index e39a1b7ea340a..c8faf5a874f5f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/ProcedureSuite.scala @@ -258,18 +258,20 @@ class ProcedureSuite extends QueryTest with SharedSparkSession with BeforeAndAft } test("malformed input to implicit cast") { - catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) - val call = "CALL cat.ns.sum('A', 2)" - checkError( - exception = intercept[SparkNumberFormatException]( - sql(call) - ), - condition = "CAST_INVALID_INPUT", - parameters = Map( - "expression" -> toSQLValue("A"), - "sourceType" -> toSQLType("STRING"), - "targetType" -> toSQLType("INT")), - context = ExpectedContext(fragment = call, start = 0, stop = call.length - 1)) + withSQLConf(SQLConf.ANSI_ENABLED.key -> true.toString) { + catalog.createProcedure(Identifier.of(Array("ns"), "sum"), UnboundSum) + val call = "CALL cat.ns.sum('A', 2)" + checkError( + exception = intercept[SparkNumberFormatException]( + sql(call) + ), + condition = "CAST_INVALID_INPUT", + parameters = Map( + "expression" -> toSQLValue("A"), + "sourceType" -> toSQLType("STRING"), + "targetType" -> toSQLType("INT")), + context = ExpectedContext(fragment = call, start = 0, stop = call.length - 1)) + } } test("required parameters after optional") { From bbbc05cbf971e931a1defc54b9924060dcdf55ca Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Sat, 21 Sep 2024 19:39:27 +0800 Subject: [PATCH 097/189] [SPARK-49495][DOCS] Document and Feature Preview on the master branch via Live GitHub Pages Updates ### What changes were proposed in this pull request? This pull request introduces functionalities that enable 'Document and Feature Preview on the master branch via Live GitHub Pages Updates'. ### Why are the changes needed? retore 8861f0f9af3f397921ba1204cf4f76f4e20680bb 376382711e200aa978008b25630cc54271fd419b 58d73fe8e7cbff9878539d31430f819eff9fc7a1 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? https://github.com/yaooqinn/spark/actions/runs/10952355999 ### Was this patch authored or co-authored using generative AI tooling? no Closes #48175 from yaooqinn/SPARK-49495. Authored-by: Kent Yao Signed-off-by: Kent Yao --- .asf.yaml | 2 + .github/workflows/pages.yml | 92 +++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 .github/workflows/pages.yml diff --git a/.asf.yaml b/.asf.yaml index 22042b355b2fa..3935a525ff3c4 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -31,6 +31,8 @@ github: merge: false squash: true rebase: true + ghp_branch: master + ghp_path: /docs notifications: pullrequests: reviews@spark.apache.org diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml new file mode 100644 index 0000000000000..b3f1cad8d947f --- /dev/null +++ b/.github/workflows/pages.yml @@ -0,0 +1,92 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +name: GitHub Pages deployment + +on: + push: + branches: + - master + +concurrency: + group: 'docs preview' + cancel-in-progress: false + +jobs: + docs: + name: Build and deploy documentation + runs-on: ubuntu-latest + permissions: + id-token: write + pages: write + environment: + name: github-pages # https://github.com/actions/deploy-pages/issues/271 + env: + SPARK_TESTING: 1 # Reduce some noise in the logs + RELEASE_VERSION: 'In-Progress' + steps: + - name: Checkout Spark repository + uses: actions/checkout@v4 + with: + repository: apache/spark + ref: 'master' + - name: Install Java 17 + uses: actions/setup-java@v4 + with: + distribution: zulu + java-version: 17 + - name: Install Python 3.9 + uses: actions/setup-python@v5 + with: + python-version: '3.9' + architecture: x64 + cache: 'pip' + - name: Install Python dependencies + run: pip install --upgrade -r dev/requirements.txt + - name: Install Ruby for documentation generation + uses: ruby/setup-ruby@v1 + with: + ruby-version: '3.3' + bundler-cache: true + - name: Install Pandoc + run: | + sudo apt-get update -y + sudo apt-get install pandoc + - name: Install dependencies for documentation generation + run: | + cd docs + gem install bundler -v 2.4.22 -n /usr/local/bin + bundle install --retry=100 + - name: Run documentation build + run: | + sed -i".tmp1" 's/SPARK_VERSION:.*$/SPARK_VERSION: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp2" 's/SPARK_VERSION_SHORT:.*$/SPARK_VERSION_SHORT: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp3" "s/'facetFilters':.*$/'facetFilters': [\"version:$RELEASE_VERSION\"]/g" docs/_config.yml + sed -i".tmp4" 's/__version__: str = .*$/__version__: str = "'"$RELEASE_VERSION"'"/' python/pyspark/version.py + cd docs + SKIP_RDOC=1 bundle exec jekyll build + - name: Setup Pages + uses: actions/configure-pages@v5 + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: 'docs/_site' + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 From 19906468d145a52a0f039e49fa54c558767805b2 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Sat, 21 Sep 2024 20:44:08 +0800 Subject: [PATCH 098/189] Revert "[SPARK-49495][DOCS] Document and Feature Preview on the master branch via Live GitHub Pages Updates" This reverts commit bbbc05cbf971e931a1defc54b9924060dcdf55ca. --- .asf.yaml | 2 - .github/workflows/pages.yml | 92 ------------------------------------- 2 files changed, 94 deletions(-) delete mode 100644 .github/workflows/pages.yml diff --git a/.asf.yaml b/.asf.yaml index 3935a525ff3c4..22042b355b2fa 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -31,8 +31,6 @@ github: merge: false squash: true rebase: true - ghp_branch: master - ghp_path: /docs notifications: pullrequests: reviews@spark.apache.org diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml deleted file mode 100644 index b3f1cad8d947f..0000000000000 --- a/.github/workflows/pages.yml +++ /dev/null @@ -1,92 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# - -name: GitHub Pages deployment - -on: - push: - branches: - - master - -concurrency: - group: 'docs preview' - cancel-in-progress: false - -jobs: - docs: - name: Build and deploy documentation - runs-on: ubuntu-latest - permissions: - id-token: write - pages: write - environment: - name: github-pages # https://github.com/actions/deploy-pages/issues/271 - env: - SPARK_TESTING: 1 # Reduce some noise in the logs - RELEASE_VERSION: 'In-Progress' - steps: - - name: Checkout Spark repository - uses: actions/checkout@v4 - with: - repository: apache/spark - ref: 'master' - - name: Install Java 17 - uses: actions/setup-java@v4 - with: - distribution: zulu - java-version: 17 - - name: Install Python 3.9 - uses: actions/setup-python@v5 - with: - python-version: '3.9' - architecture: x64 - cache: 'pip' - - name: Install Python dependencies - run: pip install --upgrade -r dev/requirements.txt - - name: Install Ruby for documentation generation - uses: ruby/setup-ruby@v1 - with: - ruby-version: '3.3' - bundler-cache: true - - name: Install Pandoc - run: | - sudo apt-get update -y - sudo apt-get install pandoc - - name: Install dependencies for documentation generation - run: | - cd docs - gem install bundler -v 2.4.22 -n /usr/local/bin - bundle install --retry=100 - - name: Run documentation build - run: | - sed -i".tmp1" 's/SPARK_VERSION:.*$/SPARK_VERSION: '"$RELEASE_VERSION"'/g' docs/_config.yml - sed -i".tmp2" 's/SPARK_VERSION_SHORT:.*$/SPARK_VERSION_SHORT: '"$RELEASE_VERSION"'/g' docs/_config.yml - sed -i".tmp3" "s/'facetFilters':.*$/'facetFilters': [\"version:$RELEASE_VERSION\"]/g" docs/_config.yml - sed -i".tmp4" 's/__version__: str = .*$/__version__: str = "'"$RELEASE_VERSION"'"/' python/pyspark/version.py - cd docs - SKIP_RDOC=1 bundle exec jekyll build - - name: Setup Pages - uses: actions/configure-pages@v5 - - name: Upload artifact - uses: actions/upload-pages-artifact@v3 - with: - path: 'docs/_site' - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 From 4f640e2485d24088345b3f2d894c696ef29e2923 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Sat, 21 Sep 2024 23:18:31 +0800 Subject: [PATCH 099/189] [SPARK-49495][DOCS] Document and Feature Preview on the master branch via Live GitHub Pages Updates ### What changes were proposed in this pull request? This pull request introduces functionalities that enable 'Document and Feature Preview on the master branch via Live GitHub Pages Updates'. ### Why are the changes needed? retore 8861f0f9af3f397921ba1204cf4f76f4e20680bb 376382711e200aa978008b25630cc54271fd419b 58d73fe8e7cbff9878539d31430f819eff9fc7a1 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? https://github.com/yaooqinn/spark/actions/runs/10952355999 ### Was this patch authored or co-authored using generative AI tooling? no Closes #48175 from yaooqinn/SPARK-49495. Authored-by: Kent Yao Signed-off-by: Kent Yao --- .asf.yaml | 2 + .github/workflows/pages.yml | 97 +++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 .github/workflows/pages.yml diff --git a/.asf.yaml b/.asf.yaml index 22042b355b2fa..3935a525ff3c4 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -31,6 +31,8 @@ github: merge: false squash: true rebase: true + ghp_branch: master + ghp_path: /docs notifications: pullrequests: reviews@spark.apache.org diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml new file mode 100644 index 0000000000000..8faeb0557fbfb --- /dev/null +++ b/.github/workflows/pages.yml @@ -0,0 +1,97 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +name: GitHub Pages deployment + +on: + push: + branches: + - master + +concurrency: + group: 'docs preview' + cancel-in-progress: false + +jobs: + docs: + name: Build and deploy documentation + runs-on: ubuntu-latest + permissions: + id-token: write + pages: write + environment: + name: github-pages # https://github.com/actions/deploy-pages/issues/271 + env: + SPARK_TESTING: 1 # Reduce some noise in the logs + RELEASE_VERSION: 'In-Progress' + steps: + - name: Checkout Spark repository + uses: actions/checkout@v4 + with: + repository: apache/spark + ref: 'master' + - name: Install Java 17 + uses: actions/setup-java@v4 + with: + distribution: zulu + java-version: 17 + - name: Install Python 3.9 + uses: actions/setup-python@v5 + with: + python-version: '3.9' + architecture: x64 + cache: 'pip' + - name: Install Python dependencies + run: | + pip install 'sphinx==4.5.0' mkdocs 'pydata_sphinx_theme>=0.13' sphinx-copybutton nbsphinx numpydoc jinja2 markupsafe 'pyzmq<24.0.0' \ + ipython ipython_genutils sphinx_plotly_directive 'numpy>=1.20.0' pyarrow 'pandas==2.2.2' 'plotly>=4.8' 'docutils<0.18.0' \ + 'flake8==3.9.0' 'mypy==1.8.0' 'pytest==7.1.3' 'pytest-mypy-plugins==1.9.3' 'black==23.9.1' \ + 'pandas-stubs==1.2.0.53' 'grpcio==1.62.0' 'grpcio-status==1.62.0' 'grpc-stubs==1.24.11' 'googleapis-common-protos-stubs==2.2.0' \ + 'sphinxcontrib-applehelp==1.0.4' 'sphinxcontrib-devhelp==1.0.2' 'sphinxcontrib-htmlhelp==2.0.1' 'sphinxcontrib-qthelp==1.0.3' 'sphinxcontrib-serializinghtml==1.1.5' + - name: Install Ruby for documentation generation + uses: ruby/setup-ruby@v1 + with: + ruby-version: '3.3' + bundler-cache: true + - name: Install Pandoc + run: | + sudo apt-get update -y + sudo apt-get install pandoc + - name: Install dependencies for documentation generation + run: | + cd docs + gem install bundler -v 2.4.22 -n /usr/local/bin + bundle install --retry=100 + - name: Run documentation build + run: | + sed -i".tmp1" 's/SPARK_VERSION:.*$/SPARK_VERSION: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp2" 's/SPARK_VERSION_SHORT:.*$/SPARK_VERSION_SHORT: '"$RELEASE_VERSION"'/g' docs/_config.yml + sed -i".tmp3" "s/'facetFilters':.*$/'facetFilters': [\"version:$RELEASE_VERSION\"]/g" docs/_config.yml + sed -i".tmp4" 's/__version__: str = .*$/__version__: str = "'"$RELEASE_VERSION"'"/' python/pyspark/version.py + cd docs + SKIP_RDOC=1 bundle exec jekyll build + - name: Setup Pages + uses: actions/configure-pages@v5 + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: 'docs/_site' + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 From b6420969b5df2ba1f542e020c5773d1d107734e9 Mon Sep 17 00:00:00 2001 From: Tim Lee Date: Sun, 22 Sep 2024 14:40:41 +0900 Subject: [PATCH 100/189] [SPARK-49741][DOCS] Add `spark.shuffle.accurateBlockSkewedFactor` to config docs page ### What changes were proposed in this pull request? `spark.shuffle.accurateBlockSkewedFactor` was added in Spark 3.3.0 in https://issues.apache.org/jira/browse/SPARK-36967 and is a useful shuffle configuration to prevent issues where `HighlyCompressedMapStatus` wrongly estimates the shuffle block sizes when the block size distribution is skewed, which can cause the shuffle reducer to fetch too much data and OOM. This PR adds this config to the Spark config docs page to make it discoverable. ### Why are the changes needed? To make this useful config discoverable by users and make them able to resolve shuffle fetch OOM issues themselves. ### Does this PR introduce _any_ user-facing change? Yes, this is a documentation fix. Before this PR there's no `spark.sql.adaptive.skewJoin.skewedPartitionFactor` in the `Shuffle Behavior` section on [the Configurations page](https://spark.apache.org/docs/latest/configuration.html) and now there is. ### How was this patch tested? On the IDE: image Updated: image ### Was this patch authored or co-authored using generative AI tooling? No Closes #48189 from timlee0119/add-accurate-block-skewed-factor-to-doc. Authored-by: Tim Lee Signed-off-by: Hyukjin Kwon --- .../org/apache/spark/internal/config/package.scala | 1 - docs/configuration.md | 13 +++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala index 47019c04aada2..c5646d2956aeb 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/package.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala @@ -1386,7 +1386,6 @@ package object config { private[spark] val SHUFFLE_ACCURATE_BLOCK_SKEWED_FACTOR = ConfigBuilder("spark.shuffle.accurateBlockSkewedFactor") - .internal() .doc("A shuffle block is considered as skewed and will be accurately recorded in " + "HighlyCompressedMapStatus if its size is larger than this factor multiplying " + "the median shuffle block size or SHUFFLE_ACCURATE_BLOCK_THRESHOLD. It is " + diff --git a/docs/configuration.md b/docs/configuration.md index 73d57b687ca2a..3c83ed92c1280 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1232,6 +1232,19 @@ Apart from these, the following properties are also available, and may be useful 2.2.1 + + spark.shuffle.accurateBlockSkewedFactor + -1.0 + + A shuffle block is considered as skewed and will be accurately recorded in + HighlyCompressedMapStatus if its size is larger than this factor multiplying + the median shuffle block size or spark.shuffle.accurateBlockThreshold. It is + recommended to set this parameter to be the same as + spark.sql.adaptive.skewJoin.skewedPartitionFactor. Set to -1.0 to disable this + feature by default. + + 3.3.0 + spark.shuffle.registration.timeout 5000 From 067f8f188eb22f9abe39eee0d70ad1ef73f4f644 Mon Sep 17 00:00:00 2001 From: Max Gekk Date: Sun, 22 Sep 2024 14:41:25 +0900 Subject: [PATCH 101/189] [SPARK-48355][SQL][TESTS][FOLLOWUP] Enable a SQL Scripting test in ANSI and non-ANSI modes ### What changes were proposed in this pull request? In the PR, I propose to enable the test which https://github.com/apache/spark/pull/48115 turned off, and run in the ANSI and non-ANSI modes. ### Why are the changes needed? To make this test stable, and don't depend on the default setting for ANSI mode. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? By running the modified test locally: ``` $ PYSPARK_PYTHON=python3 build/sbt "sql/testOnly org.apache.spark.sql.scripting.SqlScriptingInterpreterSuite" ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48194 from MaxGekk/enable-sqlscript-test-ansi. Authored-by: Max Gekk Signed-off-by: Hyukjin Kwon --- .../SqlScriptingInterpreterSuite.scala | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala index bc2adec5be3d5..ac190eb48d1f9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala @@ -21,6 +21,7 @@ import org.apache.spark.{SparkException, SparkNumberFormatException} import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, QueryTest, Row} import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.exceptions.SqlScriptingException +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession /** @@ -701,8 +702,7 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { verifySqlScriptResult(commands, expected) } - // This is disabled because it fails in non-ANSI mode - ignore("simple case mismatched types") { + test("simple case mismatched types") { val commands = """ |BEGIN @@ -712,18 +712,26 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession { | END CASE; |END |""".stripMargin - - checkError( - exception = intercept[SparkNumberFormatException] ( - runSqlScript(commands) - ), - condition = "CAST_INVALID_INPUT", - parameters = Map( - "expression" -> "'one'", - "sourceType" -> "\"STRING\"", - "targetType" -> "\"BIGINT\""), - context = ExpectedContext(fragment = "\"one\"", start = 23, stop = 27) - ) + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + checkError( + exception = intercept[SparkNumberFormatException]( + runSqlScript(commands) + ), + condition = "CAST_INVALID_INPUT", + parameters = Map( + "expression" -> "'one'", + "sourceType" -> "\"STRING\"", + "targetType" -> "\"BIGINT\""), + context = ExpectedContext(fragment = "\"one\"", start = 23, stop = 27)) + } + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + checkError( + exception = intercept[SqlScriptingException]( + runSqlScript(commands) + ), + condition = "BOOLEAN_STATEMENT_WITH_EMPTY_ROW", + parameters = Map("invalidStatement" -> "\"ONE\"")) + } } test("simple case compare with null") { From 719b57a32e0f36e7c425137014df2b83b7c4b029 Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Sun, 22 Sep 2024 14:29:24 -0700 Subject: [PATCH 102/189] [SPARK-49752][YARN] Remove workaround for YARN-3350 ### What changes were proposed in this pull request? Remove the logic of forcibly setting the log level to WARN for `org.apache.hadoop.yarn.util.RackResolver`. ### Why are the changes needed? The removed code was introduced in SPARK-5393 as a workaround for YARN-3350, which is already fixed on the YARN 2.8.0/3.0.0. ### Does this PR introduce _any_ user-facing change? Yes, previously, the log level of RackResolver is hardcoded as WARN even if the user explicitly sets it to DEBUG. ### How was this patch tested? Review. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48201 from pan3793/SPARK-49752. Authored-by: Cheng Pan Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/deploy/yarn/SparkRackResolver.scala | 9 --------- 1 file changed, 9 deletions(-) diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/SparkRackResolver.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/SparkRackResolver.scala index 618f0dc8a4daa..d6e814f5c30a5 100644 --- a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/SparkRackResolver.scala +++ b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/SparkRackResolver.scala @@ -25,9 +25,6 @@ import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.CommonConfigurationKeysPublic import org.apache.hadoop.net._ import org.apache.hadoop.util.ReflectionUtils -import org.apache.hadoop.yarn.util.RackResolver -import org.apache.logging.log4j.{Level, LogManager} -import org.apache.logging.log4j.core.Logger import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys.NODE_LOCATION @@ -39,12 +36,6 @@ import org.apache.spark.internal.LogKeys.NODE_LOCATION */ private[spark] class SparkRackResolver(conf: Configuration) extends Logging { - // RackResolver logs an INFO message whenever it resolves a rack, which is way too often. - val logger = LogManager.getLogger(classOf[RackResolver]) - if (logger.getLevel != Level.WARN) { - logger.asInstanceOf[Logger].setLevel(Level.WARN) - } - private val dnsToSwitchMapping: DNSToSwitchMapping = { val dnsToSwitchMappingClass = conf.getClass(CommonConfigurationKeysPublic.NET_TOPOLOGY_NODE_SWITCH_MAPPING_IMPL_KEY, From 0eeb61fb64e0c499610c7b9a84f9e41e923251e8 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Mon, 23 Sep 2024 10:46:08 +0800 Subject: [PATCH 103/189] [SPARK-49734][PYTHON] Add `seed` argument for function `shuffle` ### What changes were proposed in this pull request? 1, Add `seed` argument for function `shuffle`; 2, Rewrite and enable the doctest by specify the seed and control the partitioning; ### Why are the changes needed? feature parity, seed is support in SQL side ### Does this PR introduce _any_ user-facing change? yes, new argument ### How was this patch tested? updated doctest ### Was this patch authored or co-authored using generative AI tooling? no Closes #48184 from zhengruifeng/py_func_shuffle. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- .../pyspark/sql/connect/functions/builtin.py | 10 +-- python/pyspark/sql/functions/builtin.py | 69 ++++++++++--------- .../org/apache/spark/sql/functions.scala | 13 +++- 3 files changed, 53 insertions(+), 39 deletions(-) diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index 7fed175cbc8ea..2a39bc6bfddda 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -65,7 +65,6 @@ from pyspark.sql.types import ( _from_numpy_type, DataType, - LongType, StructType, ArrayType, StringType, @@ -2206,12 +2205,9 @@ def schema_of_xml(xml: Union[str, Column], options: Optional[Mapping[str, str]] schema_of_xml.__doc__ = pysparkfuncs.schema_of_xml.__doc__ -def shuffle(col: "ColumnOrName") -> Column: - return _invoke_function( - "shuffle", - _to_col(col), - LiteralExpression(random.randint(0, sys.maxsize), LongType()), - ) +def shuffle(col: "ColumnOrName", seed: Optional[Union[Column, int]] = None) -> Column: + _seed = lit(random.randint(0, sys.maxsize)) if seed is None else lit(seed) + return _invoke_function("shuffle", _to_col(col), _seed) shuffle.__doc__ = pysparkfuncs.shuffle.__doc__ diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 5f8d1c21a24f1..2d5dbb5946050 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -17723,7 +17723,7 @@ def array_sort( @_try_remote_functions -def shuffle(col: "ColumnOrName") -> Column: +def shuffle(col: "ColumnOrName", seed: Optional[Union[Column, int]] = None) -> Column: """ Array function: Generates a random permutation of the given array. @@ -17736,6 +17736,10 @@ def shuffle(col: "ColumnOrName") -> Column: ---------- col : :class:`~pyspark.sql.Column` or str The name of the column or expression to be shuffled. + seed : :class:`~pyspark.sql.Column` or int, optional + Seed value for the random generator. + + .. versionadded:: 4.0.0 Returns ------- @@ -17752,48 +17756,51 @@ def shuffle(col: "ColumnOrName") -> Column: Example 1: Shuffling a simple array >>> import pyspark.sql.functions as sf - >>> df = spark.createDataFrame([([1, 20, 3, 5],)], ['data']) - >>> df.select(sf.shuffle(df.data)).show() # doctest: +SKIP - +-------------+ - |shuffle(data)| - +-------------+ - |[1, 3, 20, 5]| - +-------------+ + >>> df = spark.sql("SELECT ARRAY(1, 20, 3, 5) AS data") + >>> df.select("*", sf.shuffle(df.data, sf.lit(123))).show() + +-------------+-------------+ + | data|shuffle(data)| + +-------------+-------------+ + |[1, 20, 3, 5]|[5, 1, 20, 3]| + +-------------+-------------+ Example 2: Shuffling an array with null values >>> import pyspark.sql.functions as sf - >>> df = spark.createDataFrame([([1, 20, None, 3],)], ['data']) - >>> df.select(sf.shuffle(df.data)).show() # doctest: +SKIP - +----------------+ - | shuffle(data)| - +----------------+ - |[20, 3, NULL, 1]| - +----------------+ + >>> df = spark.sql("SELECT ARRAY(1, 20, NULL, 5) AS data") + >>> df.select("*", sf.shuffle(sf.col("data"), 234)).show() + +----------------+----------------+ + | data| shuffle(data)| + +----------------+----------------+ + |[1, 20, NULL, 5]|[NULL, 5, 20, 1]| + +----------------+----------------+ Example 3: Shuffling an array with duplicate values >>> import pyspark.sql.functions as sf - >>> df = spark.createDataFrame([([1, 2, 2, 3, 3, 3],)], ['data']) - >>> df.select(sf.shuffle(df.data)).show() # doctest: +SKIP - +------------------+ - | shuffle(data)| - +------------------+ - |[3, 2, 1, 3, 2, 3]| - +------------------+ + >>> df = spark.sql("SELECT ARRAY(1, 2, 2, 3, 3, 3) AS data") + >>> df.select("*", sf.shuffle("data", 345)).show() + +------------------+------------------+ + | data| shuffle(data)| + +------------------+------------------+ + |[1, 2, 2, 3, 3, 3]|[2, 3, 3, 1, 2, 3]| + +------------------+------------------+ - Example 4: Shuffling an array with different types of elements + Example 4: Shuffling an array with random seed >>> import pyspark.sql.functions as sf - >>> df = spark.createDataFrame([(['a', 'b', 'c', 1, 2, 3],)], ['data']) - >>> df.select(sf.shuffle(df.data)).show() # doctest: +SKIP - +------------------+ - | shuffle(data)| - +------------------+ - |[1, c, 2, a, b, 3]| - +------------------+ + >>> df = spark.sql("SELECT ARRAY(1, 2, 2, 3, 3, 3) AS data") + >>> df.select("*", sf.shuffle("data")).show() # doctest: +SKIP + +------------------+------------------+ + | data| shuffle(data)| + +------------------+------------------+ + |[1, 2, 2, 3, 3, 3]|[3, 3, 2, 3, 2, 1]| + +------------------+------------------+ """ - return _invoke_function_over_columns("shuffle", col) + if seed is not None: + return _invoke_function_over_columns("shuffle", col, lit(seed)) + else: + return _invoke_function_over_columns("shuffle", col) @_try_remote_functions diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 0662b8f2b271f..d9bceabe88f8f 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -7252,7 +7252,18 @@ object functions { * @group array_funcs * @since 2.4.0 */ - def shuffle(e: Column): Column = Column.fn("shuffle", e, lit(SparkClassUtils.random.nextLong)) + def shuffle(e: Column): Column = shuffle(e, lit(SparkClassUtils.random.nextLong)) + + /** + * Returns a random permutation of the given array. + * + * @note + * The function is non-deterministic. + * + * @group array_funcs + * @since 4.0.0 + */ + def shuffle(e: Column, seed: Column): Column = Column.fn("shuffle", e, seed) /** * Returns a reversed string or an array with reverse order of elements. From 3c81f076ab9c72514cfc8372edd16e6da7c151d6 Mon Sep 17 00:00:00 2001 From: Andrey Gubichev Date: Mon, 23 Sep 2024 10:58:13 +0800 Subject: [PATCH 104/189] [SPARK-49653][SQL] Single join for correlated scalar subqueries ### What changes were proposed in this pull request? Single join is a left outer join that checks that there is at most 1 build row for every probe row. This PR adds single join implementation to support correlated scalar subqueries where the optimizer can't guarantee that 1 row is coming from them, e.g.: select *, (select t1.x from t1 where t1.y >= t_outer.y) from t_outer. -- this subquery is going to be rewritten as a single join that makes sure there is at most 1 matching build row for every probe row. It will issue a spark runtime error otherwise. Design doc: https://docs.google.com/document/d/1NTsvtBTB9XvvyRvH62QzWIZuw4hXktALUG1fBP7ha1Q/edit The optimizer introduces a single join in cases that were previously returning incorrect results (or were unsupported). Only hash-based implementation is supported, the optimizer makes sure we don't plan a single join as a sort-merge join. ### Why are the changes needed? Expands our subquery coverage. ### Does this PR introduce _any_ user-facing change? Yes, previously unsupported scalar subqueries should now work. ### How was this patch tested? Unit tests for the single join operator. Query tests for the subqueries. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48145 from agubichev/single_join. Authored-by: Andrey Gubichev Signed-off-by: Wenchen Fan --- .../sql/catalyst/analysis/Analyzer.scala | 2 +- .../sql/catalyst/analysis/CheckAnalysis.scala | 36 ++-- .../sql/catalyst/expressions/subquery.scala | 22 +- .../sql/catalyst/optimizer/Optimizer.scala | 9 +- .../sql/catalyst/optimizer/expressions.scala | 4 +- .../spark/sql/catalyst/optimizer/joins.scala | 8 +- .../sql/catalyst/optimizer/subquery.scala | 50 ++++- .../spark/sql/catalyst/plans/joinTypes.scala | 4 + .../plans/logical/basicLogicalOperators.scala | 10 +- .../sql/errors/QueryExecutionErrors.scala | 6 + .../apache/spark/sql/internal/SQLConf.scala | 9 + .../spark/sql/execution/SparkStrategies.scala | 11 +- .../adaptive/PlanAdaptiveSubqueries.scala | 2 +- .../joins/BroadcastNestedLoopJoinExec.scala | 44 +++- .../spark/sql/execution/joins/HashJoin.scala | 29 ++- .../sql/execution/joins/ShuffledJoin.scala | 6 +- .../scalar-subquery-group-by.sql.out | 111 ++++++++-- .../scalar-subquery-predicate.sql.out | 18 ++ .../scalar-subquery-group-by.sql | 11 +- .../scalar-subquery-predicate.sql | 3 + .../scalar-subquery-group-by.sql.out | 83 +++++-- .../scalar-subquery-predicate.sql.out | 8 + .../spark/sql/LateralColumnAliasSuite.scala | 11 +- .../org/apache/spark/sql/SubquerySuite.scala | 44 ++-- .../sql/execution/joins/SingleJoinSuite.scala | 204 ++++++++++++++++++ 25 files changed, 613 insertions(+), 132 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SingleJoinSuite.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index 9e5b1d1254c87..b2e9115dd512f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala @@ -2716,7 +2716,7 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor */ private def resolveSubQueries(plan: LogicalPlan, outer: LogicalPlan): LogicalPlan = { plan.transformAllExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION), ruleId) { - case s @ ScalarSubquery(sub, _, exprId, _, _, _) if !sub.resolved => + case s @ ScalarSubquery(sub, _, exprId, _, _, _, _) if !sub.resolved => resolveSubQuery(s, outer)(ScalarSubquery(_, _, exprId)) case e @ Exists(sub, _, exprId, _, _) if !sub.resolved => resolveSubQuery(e, outer)(Exists(_, _, exprId)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala index 5a9d5cd87ecc7..b600f455f16ac 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala @@ -952,19 +952,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB messageParameters = Map.empty) } - // SPARK-18504/SPARK-18814: Block cases where GROUP BY columns - // are not part of the correlated columns. - - // Collect the inner query expressions that are guaranteed to have a single value for each - // outer row. See comment on getCorrelatedEquivalentInnerExpressions. - val correlatedEquivalentExprs = getCorrelatedEquivalentInnerExpressions(query) - // Grouping expressions, except outer refs and constant expressions - grouping by an - // outer ref or a constant is always ok - val groupByExprs = - ExpressionSet(agg.groupingExpressions.filter(x => !x.isInstanceOf[OuterReference] && - x.references.nonEmpty)) - val nonEquivalentGroupByExprs = groupByExprs -- correlatedEquivalentExprs - + val nonEquivalentGroupByExprs = nonEquivalentGroupbyCols(query, agg) val invalidCols = if (!SQLConf.get.getConf( SQLConf.LEGACY_SCALAR_SUBQUERY_ALLOW_GROUP_BY_NON_EQUALITY_CORRELATED_PREDICATE)) { nonEquivalentGroupByExprs @@ -1044,7 +1032,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB checkOuterReference(plan, expr) expr match { - case ScalarSubquery(query, outerAttrs, _, _, _, _) => + case ScalarSubquery(query, outerAttrs, _, _, _, _, _) => // Scalar subquery must return one column as output. if (query.output.size != 1) { throw QueryCompilationErrors.subqueryReturnMoreThanOneColumn(query.output.size, @@ -1052,15 +1040,17 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB } if (outerAttrs.nonEmpty) { - cleanQueryInScalarSubquery(query) match { - case a: Aggregate => checkAggregateInScalarSubquery(outerAttrs, query, a) - case Filter(_, a: Aggregate) => checkAggregateInScalarSubquery(outerAttrs, query, a) - case p: LogicalPlan if p.maxRows.exists(_ <= 1) => // Ok - case other => - expr.failAnalysis( - errorClass = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY." + - "MUST_AGGREGATE_CORRELATED_SCALAR_SUBQUERY", - messageParameters = Map.empty) + if (!SQLConf.get.getConf(SQLConf.SCALAR_SUBQUERY_USE_SINGLE_JOIN)) { + cleanQueryInScalarSubquery(query) match { + case a: Aggregate => checkAggregateInScalarSubquery(outerAttrs, query, a) + case Filter(_, a: Aggregate) => checkAggregateInScalarSubquery(outerAttrs, query, a) + case p: LogicalPlan if p.maxRows.exists(_ <= 1) => // Ok + case other => + expr.failAnalysis( + errorClass = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY." + + "MUST_AGGREGATE_CORRELATED_SCALAR_SUBQUERY", + messageParameters = Map.empty) + } } // Only certain operators are allowed to host subquery expression containing diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/subquery.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/subquery.scala index 174d32c73fc01..0c8253659dd56 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/subquery.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/subquery.scala @@ -358,6 +358,20 @@ object SubExprUtils extends PredicateHelper { case _ => ExpressionSet().empty } } + + // Returns grouping expressions of 'aggNode' of a scalar subquery that do not have equivalent + // columns in the outer query (bound by equality predicates like 'col = outer(c)'). + // We use it to analyze whether a scalar subquery is guaranteed to return at most 1 row. + def nonEquivalentGroupbyCols(query: LogicalPlan, aggNode: Aggregate): ExpressionSet = { + val correlatedEquivalentExprs = getCorrelatedEquivalentInnerExpressions(query) + // Grouping expressions, except outer refs and constant expressions - grouping by an + // outer ref or a constant is always ok + val groupByExprs = + ExpressionSet(aggNode.groupingExpressions.filter(x => !x.isInstanceOf[OuterReference] && + x.references.nonEmpty)) + val nonEquivalentGroupByExprs = groupByExprs -- correlatedEquivalentExprs + nonEquivalentGroupByExprs + } } /** @@ -371,6 +385,11 @@ object SubExprUtils extends PredicateHelper { * case the subquery yields no row at all on empty input to the GROUP BY, which evaluates to NULL. * It is set in PullupCorrelatedPredicates to true/false, before it is set its value is None. * See constructLeftJoins in RewriteCorrelatedScalarSubquery for more details. + * + * 'needSingleJoin' is set to true if we can't guarantee that the correlated scalar subquery + * returns at most 1 row. For such subqueries we use a modification of an outer join called + * LeftSingle join. This value is set in PullupCorrelatedPredicates and used in + * RewriteCorrelatedScalarSubquery. */ case class ScalarSubquery( plan: LogicalPlan, @@ -378,7 +397,8 @@ case class ScalarSubquery( exprId: ExprId = NamedExpression.newExprId, joinCond: Seq[Expression] = Seq.empty, hint: Option[HintInfo] = None, - mayHaveCountBug: Option[Boolean] = None) + mayHaveCountBug: Option[Boolean] = None, + needSingleJoin: Option[Boolean] = None) extends SubqueryExpression(plan, outerAttrs, exprId, joinCond, hint) with Unevaluable { override def dataType: DataType = { if (!plan.schema.fields.nonEmpty) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 8e14537c6a5b4..7fc12f7d1fc16 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -338,7 +338,7 @@ abstract class Optimizer(catalogManager: CatalogManager) case d: DynamicPruningSubquery => d case s @ ScalarSubquery( PhysicalOperation(projections, predicates, a @ Aggregate(group, _, child)), - _, _, _, _, mayHaveCountBug) + _, _, _, _, mayHaveCountBug, _) if conf.getConf(SQLConf.DECORRELATE_SUBQUERY_PREVENT_CONSTANT_FOLDING_FOR_COUNT_BUG) && mayHaveCountBug.nonEmpty && mayHaveCountBug.get => // This is a subquery with an aggregate that may suffer from a COUNT bug. @@ -1988,7 +1988,8 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper { } private def canPushThrough(joinType: JoinType): Boolean = joinType match { - case _: InnerLike | LeftSemi | RightOuter | LeftOuter | LeftAnti | ExistenceJoin(_) => true + case _: InnerLike | LeftSemi | RightOuter | LeftOuter | LeftSingle | + LeftAnti | ExistenceJoin(_) => true case _ => false } @@ -2028,7 +2029,7 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper { (leftFilterConditions ++ commonFilterCondition). reduceLeftOption(And).map(Filter(_, newJoin)).getOrElse(newJoin) - case LeftOuter | LeftExistence(_) => + case LeftOuter | LeftSingle | LeftExistence(_) => // push down the left side only `where` condition val newLeft = leftFilterConditions. reduceLeftOption(And).map(Filter(_, left)).getOrElse(left) @@ -2074,6 +2075,8 @@ object PushPredicateThroughJoin extends Rule[LogicalPlan] with PredicateHelper { val newJoinCond = (leftJoinConditions ++ commonJoinCondition).reduceLeftOption(And) Join(newLeft, newRight, joinType, newJoinCond, hint) + // Do not move join predicates of a single join. + case LeftSingle => j case other => throw SparkException.internalError(s"Unexpected join type: $other") diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala index 3cdde622d51f7..1601d798283c9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala @@ -90,7 +90,7 @@ object ConstantFolding extends Rule[LogicalPlan] { } // Don't replace ScalarSubquery if its plan is an aggregate that may suffer from a COUNT bug. - case s @ ScalarSubquery(_, _, _, _, _, mayHaveCountBug) + case s @ ScalarSubquery(_, _, _, _, _, mayHaveCountBug, _) if conf.getConf(SQLConf.DECORRELATE_SUBQUERY_PREVENT_CONSTANT_FOLDING_FOR_COUNT_BUG) && mayHaveCountBug.nonEmpty && mayHaveCountBug.get => s @@ -1007,7 +1007,7 @@ object FoldablePropagation extends Rule[LogicalPlan] { replaceFoldable(j.withNewChildren(newChildren).asInstanceOf[Join], foldableMap) val missDerivedAttrsSet: AttributeSet = AttributeSet(newJoin.joinType match { case _: InnerLike | LeftExistence(_) => Nil - case LeftOuter => newJoin.right.output + case LeftOuter | LeftSingle => newJoin.right.output case RightOuter => newJoin.left.output case FullOuter => newJoin.left.output ++ newJoin.right.output case _ => Nil diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala index 9fc4873c248b5..6802adaa2ea24 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala @@ -339,8 +339,8 @@ trait JoinSelectionHelper extends Logging { ) } - def getBroadcastNestedLoopJoinBuildSide(hint: JoinHint): Option[BuildSide] = { - if (hintToNotBroadcastAndReplicateLeft(hint)) { + def getBroadcastNestedLoopJoinBuildSide(hint: JoinHint, joinType: JoinType): Option[BuildSide] = { + if (hintToNotBroadcastAndReplicateLeft(hint) || joinType == LeftSingle) { Some(BuildRight) } else if (hintToNotBroadcastAndReplicateRight(hint)) { Some(BuildLeft) @@ -375,7 +375,7 @@ trait JoinSelectionHelper extends Logging { def canBuildBroadcastRight(joinType: JoinType): Boolean = { joinType match { - case _: InnerLike | LeftOuter | LeftSemi | LeftAnti | _: ExistenceJoin => true + case _: InnerLike | LeftOuter | LeftSingle | LeftSemi | LeftAnti | _: ExistenceJoin => true case _ => false } } @@ -389,7 +389,7 @@ trait JoinSelectionHelper extends Logging { def canBuildShuffledHashJoinRight(joinType: JoinType): Boolean = { joinType match { - case _: InnerLike | LeftOuter | FullOuter | RightOuter | + case _: InnerLike | LeftOuter | LeftSingle | FullOuter | RightOuter | LeftSemi | LeftAnti | _: ExistenceJoin => true case _ => false } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala index 1239a5dde1302..d9795cf338279 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala @@ -456,6 +456,31 @@ object PullupCorrelatedPredicates extends Rule[LogicalPlan] with PredicateHelper (newPlan, newCond) } + // Returns true if 'query' is guaranteed to return at most 1 row. + private def guaranteedToReturnOneRow(query: LogicalPlan): Boolean = { + if (query.maxRows.exists(_ <= 1)) { + return true + } + val aggNode = query match { + case havingPart@Filter(_, aggPart: Aggregate) => Some(aggPart) + case aggPart: Aggregate => Some(aggPart) + // LIMIT 1 is handled above, this is for all other types of LIMITs + case Limit(_, aggPart: Aggregate) => Some(aggPart) + case Project(_, aggPart: Aggregate) => Some(aggPart) + case _: LogicalPlan => None + } + if (!aggNode.isDefined) { + return false + } + val aggregates = aggNode.get.expressions.flatMap(_.collect { + case a: AggregateExpression => a + }) + if (aggregates.isEmpty) { + return false + } + nonEquivalentGroupbyCols(query, aggNode.get).isEmpty + } + private def rewriteSubQueries(plan: LogicalPlan): LogicalPlan = { /** * This function is used as a aid to enforce idempotency of pullUpCorrelatedPredicate rule. @@ -481,7 +506,8 @@ object PullupCorrelatedPredicates extends Rule[LogicalPlan] with PredicateHelper } plan.transformExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION)) { - case ScalarSubquery(sub, children, exprId, conditions, hint, mayHaveCountBugOld) + case ScalarSubquery(sub, children, exprId, conditions, hint, + mayHaveCountBugOld, needSingleJoinOld) if children.nonEmpty => def mayHaveCountBugAgg(a: Aggregate): Boolean = { @@ -527,8 +553,13 @@ object PullupCorrelatedPredicates extends Rule[LogicalPlan] with PredicateHelper val (topPart, havingNode, aggNode) = splitSubquery(sub) (aggNode.isDefined && aggNode.get.groupingExpressions.isEmpty) } + val needSingleJoin = if (needSingleJoinOld.isDefined) { + needSingleJoinOld.get + } else { + conf.getConf(SQLConf.SCALAR_SUBQUERY_USE_SINGLE_JOIN) && !guaranteedToReturnOneRow(sub) + } ScalarSubquery(newPlan, children, exprId, getJoinCondition(newCond, conditions), - hint, Some(mayHaveCountBug)) + hint, Some(mayHaveCountBug), Some(needSingleJoin)) case Exists(sub, children, exprId, conditions, hint) if children.nonEmpty => val (newPlan, newCond) = if (SQLConf.get.decorrelateInnerQueryEnabledForExistsIn) { decorrelate(sub, plan, handleCountBug = true) @@ -786,7 +817,8 @@ object RewriteCorrelatedScalarSubquery extends Rule[LogicalPlan] with AliasHelpe subqueries: ArrayBuffer[ScalarSubquery]): (LogicalPlan, AttributeMap[Attribute]) = { val subqueryAttrMapping = ArrayBuffer[(Attribute, Attribute)]() val newChild = subqueries.foldLeft(child) { - case (currentChild, ScalarSubquery(sub, _, _, conditions, subHint, mayHaveCountBug)) => + case (currentChild, ScalarSubquery(sub, _, _, conditions, subHint, mayHaveCountBug, + needSingleJoin)) => val query = DecorrelateInnerQuery.rewriteDomainJoins(currentChild, sub, conditions) val origOutput = query.output.head // The subquery appears on the right side of the join, hence add its hint to the right @@ -794,9 +826,13 @@ object RewriteCorrelatedScalarSubquery extends Rule[LogicalPlan] with AliasHelpe val joinHint = JoinHint(None, subHint) val resultWithZeroTups = evalSubqueryOnZeroTups(query) + val joinType = needSingleJoin match { + case Some(true) => LeftSingle + case _ => LeftOuter + } lazy val planWithoutCountBug = Project( currentChild.output :+ origOutput, - Join(currentChild, query, LeftOuter, conditions.reduceOption(And), joinHint)) + Join(currentChild, query, joinType, conditions.reduceOption(And), joinHint)) if (Utils.isTesting) { assert(mayHaveCountBug.isDefined) @@ -845,7 +881,7 @@ object RewriteCorrelatedScalarSubquery extends Rule[LogicalPlan] with AliasHelpe currentChild.output :+ subqueryResultExpr, Join(currentChild, Project(query.output :+ alwaysTrueExpr, query), - LeftOuter, conditions.reduceOption(And), joinHint)) + joinType, conditions.reduceOption(And), joinHint)) } else { // CASE 3: Subquery with HAVING clause. Pull the HAVING clause above the join. @@ -877,7 +913,7 @@ object RewriteCorrelatedScalarSubquery extends Rule[LogicalPlan] with AliasHelpe currentChild.output :+ caseExpr, Join(currentChild, Project(subqueryRoot.output :+ alwaysTrueExpr, subqueryRoot), - LeftOuter, conditions.reduceOption(And), joinHint)) + joinType, conditions.reduceOption(And), joinHint)) } } } @@ -1028,7 +1064,7 @@ object OptimizeOneRowRelationSubquery extends Rule[LogicalPlan] { case p: LogicalPlan => p.transformExpressionsUpWithPruning( _.containsPattern(SCALAR_SUBQUERY)) { - case s @ ScalarSubquery(OneRowSubquery(p @ Project(_, _: OneRowRelation)), _, _, _, _, _) + case s @ ScalarSubquery(OneRowSubquery(p @ Project(_, _: OneRowRelation)), _, _, _, _, _, _) if !hasCorrelatedSubquery(s.plan) && s.joinCond.isEmpty => assert(p.projectList.size == 1) stripOuterReferences(p.projectList).head diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/joinTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/joinTypes.scala index d9da255eccc9d..41bba99673a2b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/joinTypes.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/joinTypes.scala @@ -95,6 +95,10 @@ case object LeftAnti extends JoinType { override def sql: String = "LEFT ANTI" } +case object LeftSingle extends JoinType { + override def sql: String = "LEFT SINGLE" +} + case class ExistenceJoin(exists: Attribute) extends JoinType { override def sql: String = { // This join type is only used in the end of optimizer and physical plans, we will not diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala index 90af6333b2e0b..7c549a32aca0f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala @@ -559,12 +559,12 @@ case class Join( override def maxRows: Option[Long] = { joinType match { - case Inner | Cross | FullOuter | LeftOuter | RightOuter + case Inner | Cross | FullOuter | LeftOuter | RightOuter | LeftSingle if left.maxRows.isDefined && right.maxRows.isDefined => val leftMaxRows = BigInt(left.maxRows.get) val rightMaxRows = BigInt(right.maxRows.get) val minRows = joinType match { - case LeftOuter => leftMaxRows + case LeftOuter | LeftSingle => leftMaxRows case RightOuter => rightMaxRows case FullOuter => leftMaxRows + rightMaxRows case _ => BigInt(0) @@ -590,7 +590,7 @@ case class Join( left.output :+ j.exists case LeftExistence(_) => left.output - case LeftOuter => + case LeftOuter | LeftSingle => left.output ++ right.output.map(_.withNullability(true)) case RightOuter => left.output.map(_.withNullability(true)) ++ right.output @@ -627,7 +627,7 @@ case class Join( left.constraints.union(right.constraints) case LeftExistence(_) => left.constraints - case LeftOuter => + case LeftOuter | LeftSingle => left.constraints case RightOuter => right.constraints @@ -659,7 +659,7 @@ case class Join( var patterns = Seq(JOIN) joinType match { case _: InnerLike => patterns = patterns :+ INNER_LIKE_JOIN - case LeftOuter | FullOuter | RightOuter => patterns = patterns :+ OUTER_JOIN + case LeftOuter | FullOuter | RightOuter | LeftSingle => patterns = patterns :+ OUTER_JOIN case LeftSemiOrAnti(_) => patterns = patterns :+ LEFT_SEMI_OR_ANTI_JOIN case NaturalJoin(_) | UsingJoin(_, _) => patterns = patterns :+ NATURAL_LIKE_JOIN case _ => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 4bc071155012b..4a23e9766fc5d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -2477,6 +2477,12 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE summary = getSummary(context)) } + def scalarSubqueryReturnsMultipleRows(): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "SCALAR_SUBQUERY_TOO_MANY_ROWS", + messageParameters = Map.empty) + } + def comparatorReturnsNull(firstValue: String, secondValue: String): Throwable = { new SparkException( errorClass = "COMPARATOR_RETURNS_NULL", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 6c3e9bac1cfe5..4d0930212b373 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -5090,6 +5090,15 @@ object SQLConf { .booleanConf .createWithDefault(true) + val SCALAR_SUBQUERY_USE_SINGLE_JOIN = + buildConf("spark.sql.optimizer.scalarSubqueryUseSingleJoin") + .internal() + .doc("When set to true, use LEFT_SINGLE join for correlated scalar subqueries where " + + "optimizer can't prove that only 1 row will be returned") + .version("4.0.0") + .booleanConf + .createWithDefault(true) + val ALLOW_SUBQUERY_EXPRESSIONS_IN_LAMBDAS_AND_HIGHER_ORDER_FUNCTIONS = buildConf("spark.sql.analyzer.allowSubqueryExpressionsInLambdasOrHigherOrderFunctions") .internal() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala index aee735e48fc5c..53c335c1eced6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala @@ -269,8 +269,13 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { } } + def canMerge(joinType: JoinType): Boolean = joinType match { + case LeftSingle => false + case _ => true + } + def createSortMergeJoin() = { - if (RowOrdering.isOrderable(leftKeys)) { + if (canMerge(joinType) && RowOrdering.isOrderable(leftKeys)) { Some(Seq(joins.SortMergeJoinExec( leftKeys, rightKeys, joinType, nonEquiCond, planLater(left), planLater(right)))) } else { @@ -297,7 +302,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { // This join could be very slow or OOM // Build the smaller side unless the join requires a particular build side // (e.g. NO_BROADCAST_AND_REPLICATION hint) - val requiredBuildSide = getBroadcastNestedLoopJoinBuildSide(hint) + val requiredBuildSide = getBroadcastNestedLoopJoinBuildSide(hint, joinType) val buildSide = requiredBuildSide.getOrElse(getSmallerSide(left, right)) Seq(joins.BroadcastNestedLoopJoinExec( planLater(left), planLater(right), buildSide, joinType, j.condition)) @@ -390,7 +395,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { // This join could be very slow or OOM // Build the desired side unless the join requires a particular build side // (e.g. NO_BROADCAST_AND_REPLICATION hint) - val requiredBuildSide = getBroadcastNestedLoopJoinBuildSide(hint) + val requiredBuildSide = getBroadcastNestedLoopJoinBuildSide(hint, joinType) val buildSide = requiredBuildSide.getOrElse(desiredBuildSide) Seq(joins.BroadcastNestedLoopJoinExec( planLater(left), planLater(right), buildSide, joinType, condition)) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala index df4d895867586..5f2638655c37c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala @@ -30,7 +30,7 @@ case class PlanAdaptiveSubqueries( def apply(plan: SparkPlan): SparkPlan = { plan.transformAllExpressionsWithPruning( _.containsAnyPattern(SCALAR_SUBQUERY, IN_SUBQUERY, DYNAMIC_PRUNING_SUBQUERY)) { - case expressions.ScalarSubquery(_, _, exprId, _, _, _) => + case expressions.ScalarSubquery(_, _, exprId, _, _, _, _) => val subquery = SubqueryExec.createForScalarSubquery( s"subquery#${exprId.id}", subqueryMap(exprId.id)) execution.ScalarSubquery(subquery, exprId) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala index 6dd41aca3a5e1..a7292ee1f8fa7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCo import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CodegenSupport, ExplainUtils, SparkPlan} import org.apache.spark.sql.execution.metric.SQLMetrics import org.apache.spark.util.ArrayImplicits._ @@ -63,13 +64,15 @@ case class BroadcastNestedLoopJoinExec( override def outputPartitioning: Partitioning = (joinType, buildSide) match { case (_: InnerLike, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) | - (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => streamed.outputPartitioning + (LeftSingle, BuildRight) | (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => + streamed.outputPartitioning case _ => super.outputPartitioning } override def outputOrdering: Seq[SortOrder] = (joinType, buildSide) match { case (_: InnerLike, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) | - (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => streamed.outputOrdering + (LeftSingle, BuildRight) | (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => + streamed.outputOrdering case _ => Nil } @@ -87,7 +90,7 @@ case class BroadcastNestedLoopJoinExec( joinType match { case _: InnerLike => left.output ++ right.output - case LeftOuter => + case LeftOuter | LeftSingle => left.output ++ right.output.map(_.withNullability(true)) case RightOuter => left.output.map(_.withNullability(true)) ++ right.output @@ -135,8 +138,14 @@ case class BroadcastNestedLoopJoinExec( * * LeftOuter with BuildRight * RightOuter with BuildLeft + * LeftSingle with BuildRight + * + * For the (LeftSingle, BuildRight) case we pass 'singleJoin' flag that + * makes sure there is at most 1 matching build row per every probe tuple. */ - private def outerJoin(relation: Broadcast[Array[InternalRow]]): RDD[InternalRow] = { + private def outerJoin( + relation: Broadcast[Array[InternalRow]], + singleJoin: Boolean = false): RDD[InternalRow] = { streamed.execute().mapPartitionsInternal { streamedIter => val buildRows = relation.value val joinedRow = new JoinedRow @@ -167,6 +176,9 @@ case class BroadcastNestedLoopJoinExec( resultRow = joinedRow(streamRow, buildRows(nextIndex)) nextIndex += 1 if (boundCondition(resultRow)) { + if (foundMatch && singleJoin) { + throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + } foundMatch = true return true } @@ -382,12 +394,18 @@ case class BroadcastNestedLoopJoinExec( innerJoin(broadcastedRelation) case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => outerJoin(broadcastedRelation) + case (LeftSingle, BuildRight) => + outerJoin(broadcastedRelation, singleJoin = true) case (LeftSemi, _) => leftExistenceJoin(broadcastedRelation, exists = true) case (LeftAnti, _) => leftExistenceJoin(broadcastedRelation, exists = false) case (_: ExistenceJoin, _) => existenceJoin(broadcastedRelation) + case (LeftSingle, BuildLeft) => + throw new IllegalArgumentException( + s"BroadcastNestedLoopJoin should not use the left side as build when " + + s"executing a LeftSingle join") case _ => /** * LeftOuter with BuildLeft @@ -410,7 +428,7 @@ case class BroadcastNestedLoopJoinExec( override def supportCodegen: Boolean = (joinType, buildSide) match { case (_: InnerLike, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) | - (LeftSemi | LeftAnti, BuildRight) => true + (LeftSemi | LeftAnti, BuildRight) | (LeftSingle, BuildRight) => true case _ => false } @@ -428,6 +446,7 @@ case class BroadcastNestedLoopJoinExec( (joinType, buildSide) match { case (_: InnerLike, _) => codegenInner(ctx, input) case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => codegenOuter(ctx, input) + case (LeftSingle, BuildRight) => codegenOuter(ctx, input) case (LeftSemi, BuildRight) => codegenLeftExistence(ctx, input, exists = true) case (LeftAnti, BuildRight) => codegenLeftExistence(ctx, input, exists = false) case _ => @@ -473,7 +492,9 @@ case class BroadcastNestedLoopJoinExec( """.stripMargin } - private def codegenOuter(ctx: CodegenContext, input: Seq[ExprCode]): String = { + private def codegenOuter( + ctx: CodegenContext, + input: Seq[ExprCode]): String = { val (buildRowArray, buildRowArrayTerm) = prepareBroadcast(ctx) val (buildRow, checkCondition, _) = getJoinCondition(ctx, input, streamed, broadcast) val buildVars = genOneSideJoinVars(ctx, buildRow, broadcast, setDefaultValue = true) @@ -494,12 +515,23 @@ case class BroadcastNestedLoopJoinExec( |${consume(ctx, resultVars)} """.stripMargin } else { + // For LeftSingle joins, generate the check on the number of matches. + val evaluateSingleCheck = if (joinType == LeftSingle) { + s""" + |if ($foundMatch) { + | throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + |} + |""".stripMargin + } else { + "" + } s""" |boolean $foundMatch = false; |for (int $arrayIndex = 0; $arrayIndex < $buildRowArrayTerm.length; $arrayIndex++) { | UnsafeRow $buildRow = (UnsafeRow) $buildRowArrayTerm[$arrayIndex]; | boolean $shouldOutputRow = false; | $checkCondition { + | $evaluateSingleCheck | $shouldOutputRow = true; | $foundMatch = true; | } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala index 5d59a48d544a0..ce7d48babc91e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CodegenSupport, ExplainUtils, RowIterator} import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.types.{BooleanType, IntegralType, LongType} @@ -52,7 +53,7 @@ trait HashJoin extends JoinCodegenSupport { joinType match { case _: InnerLike => left.output ++ right.output - case LeftOuter => + case LeftOuter | LeftSingle => left.output ++ right.output.map(_.withNullability(true)) case RightOuter => left.output.map(_.withNullability(true)) ++ right.output @@ -75,7 +76,7 @@ trait HashJoin extends JoinCodegenSupport { } case BuildRight => joinType match { - case _: InnerLike | LeftOuter | LeftSemi | LeftAnti | _: ExistenceJoin => + case _: InnerLike | LeftOuter | LeftSingle | LeftSemi | LeftAnti | _: ExistenceJoin => left.outputPartitioning case x => throw new IllegalArgumentException( @@ -93,7 +94,7 @@ trait HashJoin extends JoinCodegenSupport { } case BuildRight => joinType match { - case _: InnerLike | LeftOuter | LeftSemi | LeftAnti | _: ExistenceJoin => + case _: InnerLike | LeftOuter | LeftSingle | LeftSemi | LeftAnti | _: ExistenceJoin => left.outputOrdering case x => throw new IllegalArgumentException( @@ -191,7 +192,8 @@ trait HashJoin extends JoinCodegenSupport { private def outerJoin( streamedIter: Iterator[InternalRow], - hashedRelation: HashedRelation): Iterator[InternalRow] = { + hashedRelation: HashedRelation, + singleJoin: Boolean = false): Iterator[InternalRow] = { val joinedRow = new JoinedRow() val keyGenerator = streamSideKeyGenerator() val nullRow = new GenericInternalRow(buildPlan.output.length) @@ -218,6 +220,9 @@ trait HashJoin extends JoinCodegenSupport { while (buildIter != null && buildIter.hasNext) { val nextBuildRow = buildIter.next() if (boundCondition(joinedRow.withRight(nextBuildRow))) { + if (found && singleJoin) { + throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + } found = true return true } @@ -329,6 +334,8 @@ trait HashJoin extends JoinCodegenSupport { innerJoin(streamedIter, hashed) case LeftOuter | RightOuter => outerJoin(streamedIter, hashed) + case LeftSingle => + outerJoin(streamedIter, hashed, singleJoin = true) case LeftSemi => semiJoin(streamedIter, hashed) case LeftAnti => @@ -354,7 +361,7 @@ trait HashJoin extends JoinCodegenSupport { override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = { joinType match { case _: InnerLike => codegenInner(ctx, input) - case LeftOuter | RightOuter => codegenOuter(ctx, input) + case LeftOuter | RightOuter | LeftSingle => codegenOuter(ctx, input) case LeftSemi => codegenSemi(ctx, input) case LeftAnti => codegenAnti(ctx, input) case _: ExistenceJoin => codegenExistence(ctx, input) @@ -492,6 +499,17 @@ trait HashJoin extends JoinCodegenSupport { val matches = ctx.freshName("matches") val iteratorCls = classOf[Iterator[UnsafeRow]].getName val found = ctx.freshName("found") + // For LeftSingle joins generate the check on the number of build rows that match every + // probe row. Return an error for >1 matches. + val evaluateSingleCheck = if (joinType == LeftSingle) { + s""" + |if ($found) { + | throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + |} + |""".stripMargin + } else { + "" + } s""" |// generate join key for stream side @@ -505,6 +523,7 @@ trait HashJoin extends JoinCodegenSupport { | (UnsafeRow) $matches.next() : null; | ${checkCondition.trim} | if ($conditionPassed) { + | $evaluateSingleCheck | $found = true; | $numOutput.add(1); | ${consume(ctx, resultVars)} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala index 7c4628c8576c5..60e5a7769a503 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.joins import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, FullOuter, InnerLike, LeftExistence, LeftOuter, RightOuter} +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, FullOuter, InnerLike, LeftExistence, LeftOuter, LeftSingle, RightOuter} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution, Partitioning, PartitioningCollection, UnknownPartitioning, UnspecifiedDistribution} /** @@ -47,7 +47,7 @@ trait ShuffledJoin extends JoinCodegenSupport { override def outputPartitioning: Partitioning = joinType match { case _: InnerLike => PartitioningCollection(Seq(left.outputPartitioning, right.outputPartitioning)) - case LeftOuter => left.outputPartitioning + case LeftOuter | LeftSingle => left.outputPartitioning case RightOuter => right.outputPartitioning case FullOuter => UnknownPartitioning(left.outputPartitioning.numPartitions) case LeftExistence(_) => left.outputPartitioning @@ -60,7 +60,7 @@ trait ShuffledJoin extends JoinCodegenSupport { joinType match { case _: InnerLike => left.output ++ right.output - case LeftOuter => + case LeftOuter | LeftSingle => left.output ++ right.output.map(_.withNullability(true)) case RightOuter => left.output.map(_.withNullability(true)) ++ right.output diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out index bea91e09b0053..01de7beda551d 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out @@ -142,6 +142,12 @@ Project [x1#x, x2#x, scalar-subquery#x [x1#x && x2#x] AS scalarsubquery(x1, x2)# +- LocalRelation [col1#x, col2#x] +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false +-- !query analysis +SetCommand (spark.sql.optimizer.scalarSubqueryUseSingleJoin,Some(false)) + + -- !query select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 -- !query analysis @@ -202,24 +208,83 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = true +-- !query analysis +SetCommand (spark.sql.optimizer.scalarSubqueryUseSingleJoin,Some(true)) + + +-- !query +select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 +-- !query analysis +Project [x1#x, x2#x] ++- Filter (scalar-subquery#x [x1#x] = cast(1 as bigint)) + : +- Aggregate [y1#x], [count(1) AS count(1)#xL] + : +- Filter (y1#x > outer(x1#x)) + : +- SubqueryAlias y + : +- View (`y`, [y1#x, y2#x]) + : +- Project [cast(col1#x as int) AS y1#x, cast(col2#x as int) AS y2#x] + : +- LocalRelation [col1#x, col2#x] + +- SubqueryAlias x + +- View (`x`, [x1#x, x2#x]) + +- Project [cast(col1#x as int) AS x1#x, cast(col2#x as int) AS x2#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +select *, (select count(*) from y where y1 + y2 = x1 group by y1) from x +-- !query analysis +Project [x1#x, x2#x, scalar-subquery#x [x1#x] AS scalarsubquery(x1)#xL] +: +- Aggregate [y1#x], [count(1) AS count(1)#xL] +: +- Filter ((y1#x + y2#x) = outer(x1#x)) +: +- SubqueryAlias y +: +- View (`y`, [y1#x, y2#x]) +: +- Project [cast(col1#x as int) AS y1#x, cast(col2#x as int) AS y2#x] +: +- LocalRelation [col1#x, col2#x] ++- SubqueryAlias x + +- View (`x`, [x1#x, x2#x]) + +- Project [cast(col1#x as int) AS x1#x, cast(col2#x as int) AS x2#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +select *, (select count(*) from y where x1 = y1 and y2 + 10 = x1 + 1 group by y2) from x +-- !query analysis +Project [x1#x, x2#x, scalar-subquery#x [x1#x && x1#x] AS scalarsubquery(x1, x1)#xL] +: +- Aggregate [y2#x], [count(1) AS count(1)#xL] +: +- Filter ((outer(x1#x) = y1#x) AND ((y2#x + 10) = (outer(x1#x) + 1))) +: +- SubqueryAlias y +: +- View (`y`, [y1#x, y2#x]) +: +- Project [cast(col1#x as int) AS y1#x, cast(col2#x as int) AS y2#x] +: +- LocalRelation [col1#x, col2#x] ++- SubqueryAlias x + +- View (`x`, [x1#x, x2#x]) + +- Project [cast(col1#x as int) AS x1#x, cast(col2#x as int) AS x2#x] + +- LocalRelation [col1#x, col2#x] + + -- !query select *, (select count(*) from (select * from y where y1 = x1 union all select * from y) sub group by y1) from x -- !query analysis -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.NON_CORRELATED_COLUMNS_IN_GROUP_BY", - "sqlState" : "0A000", - "messageParameters" : { - "value" : "y1" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 11, - "stopIndex" : 106, - "fragment" : "(select count(*) from (select * from y where y1 = x1 union all select * from y) sub group by y1)" - } ] -} +Project [x1#x, x2#x, scalar-subquery#x [x1#x] AS scalarsubquery(x1)#xL] +: +- Aggregate [y1#x], [count(1) AS count(1)#xL] +: +- SubqueryAlias sub +: +- Union false, false +: :- Project [y1#x, y2#x] +: : +- Filter (y1#x = outer(x1#x)) +: : +- SubqueryAlias y +: : +- View (`y`, [y1#x, y2#x]) +: : +- Project [cast(col1#x as int) AS y1#x, cast(col2#x as int) AS y2#x] +: : +- LocalRelation [col1#x, col2#x] +: +- Project [y1#x, y2#x] +: +- SubqueryAlias y +: +- View (`y`, [y1#x, y2#x]) +: +- Project [cast(col1#x as int) AS y1#x, cast(col2#x as int) AS y2#x] +: +- LocalRelation [col1#x, col2#x] ++- SubqueryAlias x + +- View (`x`, [x1#x, x2#x]) + +- Project [cast(col1#x as int) AS x1#x, cast(col2#x as int) AS x2#x] + +- LocalRelation [col1#x, col2#x] -- !query @@ -227,17 +292,17 @@ select *, (select count(*) from y left join (select * from z where z1 = x1) sub -- !query analysis org.apache.spark.sql.catalyst.ExtendedAnalysisException { - "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.NON_CORRELATED_COLUMNS_IN_GROUP_BY", + "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.ACCESSING_OUTER_QUERY_COLUMN_IS_NOT_ALLOWED", "sqlState" : "0A000", "messageParameters" : { - "value" : "z1" + "treeNode" : "Filter (z1#x = outer(x1#x))\n+- SubqueryAlias z\n +- View (`z`, [z1#x, z2#x])\n +- Project [cast(col1#x as int) AS z1#x, cast(col2#x as int) AS z2#x]\n +- LocalRelation [col1#x, col2#x]\n" }, "queryContext" : [ { "objectType" : "", "objectName" : "", - "startIndex" : 11, - "stopIndex" : 103, - "fragment" : "(select count(*) from y left join (select * from z where z1 = x1) sub on y2 = z2 group by z1)" + "startIndex" : 46, + "stopIndex" : 74, + "fragment" : "select * from z where z1 = x1" } ] } @@ -248,6 +313,12 @@ set spark.sql.legacy.scalarSubqueryAllowGroupByNonEqualityCorrelatedPredicate = SetCommand (spark.sql.legacy.scalarSubqueryAllowGroupByNonEqualityCorrelatedPredicate,Some(true)) +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false +-- !query analysis +SetCommand (spark.sql.optimizer.scalarSubqueryUseSingleJoin,Some(false)) + + -- !query select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out index e3ce85fe5d209..4ff0222d6e965 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out @@ -1748,3 +1748,21 @@ Project [t1a#x, t1b#x, t1c#x] +- View (`t1`, [t1a#x, t1b#x, t1c#x]) +- Project [cast(col1#x as int) AS t1a#x, cast(col2#x as int) AS t1b#x, cast(col3#x as int) AS t1c#x] +- LocalRelation [col1#x, col2#x, col3#x] + + +-- !query +SELECT * FROM t0 WHERE t0a = (SELECT distinct(t1c) FROM t1 WHERE t1a = t0a) +-- !query analysis +Project [t0a#x, t0b#x] ++- Filter (t0a#x = scalar-subquery#x [t0a#x]) + : +- Distinct + : +- Project [t1c#x] + : +- Filter (t1a#x = outer(t0a#x)) + : +- SubqueryAlias t1 + : +- View (`t1`, [t1a#x, t1b#x, t1c#x]) + : +- Project [cast(col1#x as int) AS t1a#x, cast(col2#x as int) AS t1b#x, cast(col3#x as int) AS t1c#x] + : +- LocalRelation [col1#x, col2#x, col3#x] + +- SubqueryAlias t0 + +- View (`t0`, [t0a#x, t0b#x]) + +- Project [cast(col1#x as int) AS t0a#x, cast(col2#x as int) AS t0b#x] + +- LocalRelation [col1#x, col2#x] diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-group-by.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-group-by.sql index db7cdc97614cb..a23083e9e0e4d 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-group-by.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-group-by.sql @@ -22,16 +22,25 @@ select *, (select count(*) from y where x1 = y1 and cast(y2 as double) = x1 + 1 select *, (select count(*) from y where y2 + 1 = x1 + x2 group by y2 + 1) from x; --- Illegal queries +-- Illegal queries (single join disabled) +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false; select * from x where (select count(*) from y where y1 > x1 group by y1) = 1; select *, (select count(*) from y where y1 + y2 = x1 group by y1) from x; select *, (select count(*) from y where x1 = y1 and y2 + 10 = x1 + 1 group by y2) from x; +-- Same queries, with LeftSingle join +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = true; +select * from x where (select count(*) from y where y1 > x1 group by y1) = 1; +select *, (select count(*) from y where y1 + y2 = x1 group by y1) from x; +select *, (select count(*) from y where x1 = y1 and y2 + 10 = x1 + 1 group by y2) from x; + + -- Certain other operators like OUTER JOIN or UNION between the correlating filter and the group-by also can cause the scalar subquery to return multiple values and hence make the query illegal. select *, (select count(*) from (select * from y where y1 = x1 union all select * from y) sub group by y1) from x; select *, (select count(*) from y left join (select * from z where z1 = x1) sub on y2 = z2 group by z1) from x; -- The correlation below the join is unsupported in Spark anyway, but when we do support it this query should still be disallowed. -- Test legacy behavior conf set spark.sql.legacy.scalarSubqueryAllowGroupByNonEqualityCorrelatedPredicate = true; +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false; select * from x where (select count(*) from y where y1 > x1 group by y1) = 1; reset spark.sql.legacy.scalarSubqueryAllowGroupByNonEqualityCorrelatedPredicate; diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-predicate.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-predicate.sql index 2823888e6e438..81e0c5f98d82b 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-predicate.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/scalar-subquery/scalar-subquery-predicate.sql @@ -529,3 +529,6 @@ FROM t1 WHERE (SELECT max(t2c) FROM t2 WHERE t1b = t2b ) between 1 and 2; + + +SELECT * FROM t0 WHERE t0a = (SELECT distinct(t1c) FROM t1 WHERE t1a = t0a); diff --git a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out index 41cba1f43745f..56932edd4e545 100644 --- a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-group-by.sql.out @@ -112,6 +112,14 @@ struct 2 2 NULL +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false +-- !query schema +struct +-- !query output +spark.sql.optimizer.scalarSubqueryUseSingleJoin false + + -- !query select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 -- !query schema @@ -178,25 +186,56 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = true +-- !query schema +struct +-- !query output +spark.sql.optimizer.scalarSubqueryUseSingleJoin true + + +-- !query +select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "SCALAR_SUBQUERY_TOO_MANY_ROWS", + "sqlState" : "21000" +} + + +-- !query +select *, (select count(*) from y where y1 + y2 = x1 group by y1) from x +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "SCALAR_SUBQUERY_TOO_MANY_ROWS", + "sqlState" : "21000" +} + + +-- !query +select *, (select count(*) from y where x1 = y1 and y2 + 10 = x1 + 1 group by y2) from x +-- !query schema +struct +-- !query output +1 1 NULL +2 2 NULL + + -- !query select *, (select count(*) from (select * from y where y1 = x1 union all select * from y) sub group by y1) from x -- !query schema struct<> -- !query output -org.apache.spark.sql.catalyst.ExtendedAnalysisException +org.apache.spark.SparkRuntimeException { - "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.NON_CORRELATED_COLUMNS_IN_GROUP_BY", - "sqlState" : "0A000", - "messageParameters" : { - "value" : "y1" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 11, - "stopIndex" : 106, - "fragment" : "(select count(*) from (select * from y where y1 = x1 union all select * from y) sub group by y1)" - } ] + "errorClass" : "SCALAR_SUBQUERY_TOO_MANY_ROWS", + "sqlState" : "21000" } @@ -207,17 +246,17 @@ struct<> -- !query output org.apache.spark.sql.catalyst.ExtendedAnalysisException { - "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.NON_CORRELATED_COLUMNS_IN_GROUP_BY", + "errorClass" : "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.ACCESSING_OUTER_QUERY_COLUMN_IS_NOT_ALLOWED", "sqlState" : "0A000", "messageParameters" : { - "value" : "z1" + "treeNode" : "Filter (z1#x = outer(x1#x))\n+- SubqueryAlias z\n +- View (`z`, [z1#x, z2#x])\n +- Project [cast(col1#x as int) AS z1#x, cast(col2#x as int) AS z2#x]\n +- LocalRelation [col1#x, col2#x]\n" }, "queryContext" : [ { "objectType" : "", "objectName" : "", - "startIndex" : 11, - "stopIndex" : 103, - "fragment" : "(select count(*) from y left join (select * from z where z1 = x1) sub on y2 = z2 group by z1)" + "startIndex" : 46, + "stopIndex" : 74, + "fragment" : "select * from z where z1 = x1" } ] } @@ -230,6 +269,14 @@ struct spark.sql.legacy.scalarSubqueryAllowGroupByNonEqualityCorrelatedPredicate true +-- !query +set spark.sql.optimizer.scalarSubqueryUseSingleJoin = false +-- !query schema +struct +-- !query output +spark.sql.optimizer.scalarSubqueryUseSingleJoin false + + -- !query select * from x where (select count(*) from y where y1 > x1 group by y1) = 1 -- !query schema diff --git a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out index a02f0c70be6da..2460c2452ea56 100644 --- a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out @@ -906,3 +906,11 @@ WHERE (SELECT max(t2c) struct -- !query output + + +-- !query +SELECT * FROM t0 WHERE t0a = (SELECT distinct(t1c) FROM t1 WHERE t1a = t0a) +-- !query schema +struct +-- !query output + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/LateralColumnAliasSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/LateralColumnAliasSuite.scala index 9afba65183974..a892cd4db02b0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/LateralColumnAliasSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/LateralColumnAliasSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql import org.scalactic.source.Position import org.scalatest.Tag +import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, ExpressionSet} import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.catalyst.plans.logical.Aggregate @@ -554,7 +555,15 @@ class LateralColumnAliasSuite extends LateralColumnAliasSuiteBase { | FROM (SELECT dept * 2.0 AS id, id + 1 AS id2 FROM $testTable)) > 5 |ORDER BY id |""".stripMargin - withLCAOff { intercept[AnalysisException] { sql(query4) } } + withLCAOff { + val exception = intercept[SparkRuntimeException] { + sql(query4).collect() + } + checkError( + exception, + condition = "SCALAR_SUBQUERY_TOO_MANY_ROWS" + ) + } withLCAOn { val analyzedPlan = sql(query4).queryExecution.analyzed assert(!analyzedPlan.containsPattern(OUTER_REFERENCE)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index 23c4d51983bb4..6e160b4407ca8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql import scala.collection.mutable.ArrayBuffer +import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.catalyst.expressions.SubqueryExpression import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union} @@ -527,43 +528,30 @@ class SubquerySuite extends QueryTest test("SPARK-18504 extra GROUP BY column in correlated scalar subquery is not permitted") { withTempView("v") { Seq((1, 1), (1, 2)).toDF("c1", "c2").createOrReplaceTempView("v") - - val exception = intercept[AnalysisException] { - sql("select (select sum(-1) from v t2 where t1.c2 = t2.c1 group by t2.c2) sum from v t1") + val exception = intercept[SparkRuntimeException] { + sql("select (select sum(-1) from v t2 where t1.c2 = t2.c1 group by t2.c2) sum from v t1"). + collect() } checkError( exception, - condition = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY." + - "NON_CORRELATED_COLUMNS_IN_GROUP_BY", - parameters = Map("value" -> "c2"), - sqlState = None, - context = ExpectedContext( - fragment = "(select sum(-1) from v t2 where t1.c2 = t2.c1 group by t2.c2)", - start = 7, stop = 67)) } + condition = "SCALAR_SUBQUERY_TOO_MANY_ROWS" + ) + } } test("non-aggregated correlated scalar subquery") { - val exception1 = intercept[AnalysisException] { - sql("select a, (select b from l l2 where l2.a = l1.a) sum_b from l l1") + val exception1 = intercept[SparkRuntimeException] { + sql("select a, (select b from l l2 where l2.a = l1.a) sum_b from l l1").collect() } checkError( exception1, - condition = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY." + - "MUST_AGGREGATE_CORRELATED_SCALAR_SUBQUERY", - parameters = Map.empty, - context = ExpectedContext( - fragment = "(select b from l l2 where l2.a = l1.a)", start = 10, stop = 47)) - val exception2 = intercept[AnalysisException] { - sql("select a, (select b from l l2 where l2.a = l1.a group by 1) sum_b from l l1") - } - checkErrorMatchPVals( - exception2, - condition = "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY." + - "MUST_AGGREGATE_CORRELATED_SCALAR_SUBQUERY", - parameters = Map.empty[String, String], - sqlState = None, - context = ExpectedContext( - fragment = "(select b from l l2 where l2.a = l1.a group by 1)", start = 10, stop = 58)) + condition = "SCALAR_SUBQUERY_TOO_MANY_ROWS" + ) + checkAnswer( + sql("select a, (select b from l l2 where l2.a = l1.a group by 1) sum_b from l l1"), + Row(1, 2.0) :: Row(1, 2.0) :: Row(2, 1.0) :: Row(2, 1.0) :: Row(3, 3.0) :: + Row(null, null) :: Row(null, null) :: Row(6, null) :: Nil + ) } test("non-equal correlated scalar subquery") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SingleJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SingleJoinSuite.scala new file mode 100644 index 0000000000000..a318769af6871 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SingleJoinSuite.scala @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.joins + +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.optimizer.BuildRight +import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, Project} +import org.apache.spark.sql.execution.{SparkPlan, SparkPlanTest} +import org.apache.spark.sql.execution.exchange.EnsureRequirements +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DoubleType, IntegerType, StructType} + +class SingleJoinSuite extends SparkPlanTest with SharedSparkSession { + import testImplicits.toRichColumn + + private val EnsureRequirements = new EnsureRequirements() + + private lazy val left = spark.createDataFrame( + sparkContext.parallelize(Seq( + Row(1, 2.0), + Row(1, 2.0), + Row(2, 1.0), + Row(2, 1.0), + Row(3, 3.0), + Row(null, null), + Row(null, 5.0), + Row(6, null) + )), new StructType().add("a", IntegerType).add("b", DoubleType)) + + // (a > c && a != 6) + + private lazy val right = spark.createDataFrame( + sparkContext.parallelize(Seq( + Row(2, 3.0), + Row(3, 2.0), + Row(4, 1.0), + Row(4, 2.0), + Row(null, null), + Row(null, 5.0), + Row(6, null) + )), new StructType().add("c", IntegerType).add("d", DoubleType)) + + private lazy val singleConditionEQ = EqualTo(left.col("a").expr, right.col("c").expr) + + private lazy val nonEqualityCond = And(GreaterThan(left.col("a").expr, right.col("c").expr), + Not(EqualTo(left.col("a").expr, Literal(6)))) + + + + private def testSingleJoin( + testName: String, + leftRows: => DataFrame, + rightRows: => DataFrame, + condition: => Option[Expression], + expectedAnswer: Seq[Row], + expectError: Boolean = false): Unit = { + + def extractJoinParts(): Option[ExtractEquiJoinKeys.ReturnType] = { + val join = Join(leftRows.logicalPlan, rightRows.logicalPlan, + Inner, condition, JoinHint.NONE) + ExtractEquiJoinKeys.unapply(join) + } + + def checkSingleJoinError(planFunction: (SparkPlan, SparkPlan) => SparkPlan): Unit = { + val outputPlan = planFunction(leftRows.queryExecution.sparkPlan, + rightRows.queryExecution.sparkPlan) + checkError( + exception = intercept[SparkRuntimeException] { + SparkPlanTest.executePlan(outputPlan, spark.sqlContext) + }, + condition = "SCALAR_SUBQUERY_TOO_MANY_ROWS", + parameters = Map.empty + ) + } + + testWithWholeStageCodegenOnAndOff(s"$testName using BroadcastHashJoin") { _ => + extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => + val planFunction = (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply(BroadcastHashJoinExec( + leftKeys, rightKeys, LeftSingle, BuildRight, boundCondition, left, right)) + if (expectError) { + checkSingleJoinError(planFunction) + } else { + checkAnswer2(leftRows, rightRows, planFunction, + expectedAnswer, + sortAnswers = true) + } + } + } + testWithWholeStageCodegenOnAndOff(s"$testName using ShuffledHashJoin") { _ => + extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => + val planFunction = (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + ShuffledHashJoinExec( + leftKeys, rightKeys, LeftSingle, BuildRight, boundCondition, left, right)) + if (expectError) { + checkSingleJoinError(planFunction) + } else { + checkAnswer2(leftRows, rightRows, planFunction, + expectedAnswer, + sortAnswers = true) + } + } + } + + testWithWholeStageCodegenOnAndOff(s"$testName using BroadcastNestedLoopJoin") { _ => + val planFunction = (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + BroadcastNestedLoopJoinExec(left, right, BuildRight, LeftSingle, condition)) + if (expectError) { + checkSingleJoinError(planFunction) + } else { + checkAnswer2(leftRows, rightRows, planFunction, + expectedAnswer, + sortAnswers = true) + } + } + } + + testSingleJoin( + "test single condition (equal) for a left single join", + left, + Project(Seq(right.col("c").expr.asInstanceOf[NamedExpression]), right.logicalPlan), + Some(singleConditionEQ), + Seq(Row(1, 2.0, null), + Row(1, 2.0, null), + Row(2, 1.0, 2), + Row(2, 1.0, 2), + Row(3, 3.0, 3), + Row(6, null, 6), + Row(null, 5.0, null), + Row(null, null, null))) + + testSingleJoin( + "test single condition (equal) for a left single join -- multiple matches", + left, + Project(Seq(right.col("d").expr.asInstanceOf[NamedExpression]), right.logicalPlan), + Some(EqualTo(left.col("b").expr, right.col("d").expr)), + Seq.empty, true) + + testSingleJoin( + "test non-equality for a left single join", + left, + Project(Seq(right.col("c").expr.asInstanceOf[NamedExpression]), right.logicalPlan), + Some(nonEqualityCond), + Seq(Row(1, 2.0, null), + Row(1, 2.0, null), + Row(2, 1.0, null), + Row(2, 1.0, null), + Row(3, 3.0, 2), + Row(6, null, null), + Row(null, 5.0, null), + Row(null, null, null))) + + testSingleJoin( + "test non-equality for a left single join -- multiple matches", + left, + Project(Seq(right.col("c").expr.asInstanceOf[NamedExpression]), right.logicalPlan), + Some(GreaterThan(left.col("a").expr, right.col("c").expr)), + Seq.empty, expectError = true) + + private lazy val emptyFrame = spark.createDataFrame( + spark.sparkContext.emptyRDD[Row], new StructType().add("c", IntegerType).add("d", DoubleType)) + + testSingleJoin( + "empty inner (right) side", + left, + Project(Seq(emptyFrame.col("c").expr.asInstanceOf[NamedExpression]), emptyFrame.logicalPlan), + Some(GreaterThan(left.col("a").expr, emptyFrame.col("c").expr)), + Seq(Row(1, 2.0, null), + Row(1, 2.0, null), + Row(2, 1.0, null), + Row(2, 1.0, null), + Row(3, 3.0, null), + Row(6, null, null), + Row(null, 5.0, null), + Row(null, null, null))) + + testSingleJoin( + "empty outer (left) side", + Project(Seq(emptyFrame.col("c").expr.asInstanceOf[NamedExpression]), emptyFrame.logicalPlan), + right, + Some(EqualTo(emptyFrame.col("c").expr, right.col("c").expr)), + Seq.empty) +} From d2e8c1cb60e34a1c7e92374c07d682aa5ca79145 Mon Sep 17 00:00:00 2001 From: Julek Sompolski Date: Mon, 23 Sep 2024 12:39:02 +0900 Subject: [PATCH 105/189] [SPARK-48195][CORE] Save and reuse RDD/Broadcast created by SparkPlan ### What changes were proposed in this pull request? Save the RDD created by doExecute, instead of creating a new one in execute each time. Currently, many types of SparkPlans already save the RDD they create. For example, shuffle just save `lazy val inputRDD: RDD[InternalRow] = child.execute()`. It creates inconsistencies when an action (e.g. repeated `df.collect()`) is executed on Dataframe twice: * The SparkPlan will be reused, since the same `df.queryExecution.executedPlan` will be used. * Any not-result stage will be reused, as the shuffle operators will just have their `inputRDD` reused. * However, for result stage, `execute()` will call `doExecute()` again, and the logic of generating the actual execution RDD will be reexecuted for the result stage. This means that for example for the result stage, WSCG code gen will generate and compile new code, create a new RDD out of it. Generation of execution RDDs is also often influenced by config: for example, staying with WSCG, various configs like `spark.sql.codegen.hugeMethodLimit` or `spark.sql.codegen.methodSplitThreshold`. The fact that upon re-execution this will be evaluated anew for the result stage, but not for earlier stages creates inconsistencies in what config changes are visible. By saving the result of `doExecute` and reusing the RDD in `execute` we make sure that work in creating that RDD is not duplicated, and it is more consistent that all RDDs of the plan are reused, same as with the `executedPlan`. Note, that while the results of earlier shuffle stages are also reused, the result stage still does get executed again, as the result of it are not saved and available for Reuse in BlockManager. We also add a `Lazy` utility instead of using `lazy val` to deal with shortcomings of scala lazy val. ### Why are the changes needed? Resolved subtle inconsistencies coming from object reuse vs. recreating objects from scratch. ### Does this PR introduce _any_ user-facing change? Subtle changes caused by the RDD being reused, e.g. when a config change might be picked up. However, it makes things more consistent. Spark 4.0.0 might be a good candidate for making such a change. ### How was this patch tested? Existing SQL execution tests validate that the change in SparkPlan works. Tests were added for the new Lazy utility. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Github Copilot (trivial code completion suggestions) Closes #48037 from juliuszsompolski/SPARK-48195-rdd. Lead-authored-by: Julek Sompolski Co-authored-by: Hyukjin Kwon Co-authored-by: Wenchen Fan Signed-off-by: Hyukjin Kwon --- .../scala/org/apache/spark/util/LazyTry.scala | 70 ++++++++ .../scala/org/apache/spark/util/Utils.scala | 80 ++++++++++ .../org/apache/spark/util/LazyTrySuite.scala | 151 ++++++++++++++++++ .../org/apache/spark/util/UtilsSuite.scala | 112 ++++++++++++- .../sql/execution/CollectMetricsExec.scala | 5 + .../spark/sql/execution/SparkPlan.scala | 21 ++- .../columnar/InMemoryTableScanExec.scala | 82 +++++----- .../exchange/ShuffleExchangeExec.scala | 13 +- 8 files changed, 475 insertions(+), 59 deletions(-) create mode 100644 core/src/main/scala/org/apache/spark/util/LazyTry.scala create mode 100644 core/src/test/scala/org/apache/spark/util/LazyTrySuite.scala diff --git a/core/src/main/scala/org/apache/spark/util/LazyTry.scala b/core/src/main/scala/org/apache/spark/util/LazyTry.scala new file mode 100644 index 0000000000000..7edc08672c26b --- /dev/null +++ b/core/src/main/scala/org/apache/spark/util/LazyTry.scala @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.util + +import scala.util.Try + +/** + * Wrapper utility for a lazy val, with two differences compared to scala behavior: + * + * 1. Non-retrying in case of failure. This wrapper stores the exception in a Try, and will re-throw + * it on the access to `get`. + * In scala, when a `lazy val` field initialization throws an exception, the field remains + * uninitialized, and initialization will be re-attempted on the next access. This also can lead + * to performance issues, needlessly computing something towards a failure, and also can lead to + * duplicated side effects. + * + * 2. Resolving locking issues. + * In scala, when a `lazy val` field is initialized, it grabs the synchronized lock on the + * enclosing object instance. This can lead both to performance issues, and deadlocks. + * For example: + * a) Thread 1 entered a synchronized method, grabbing a coarse lock on the parent object. + * b) Thread 2 get spawned off, and tries to initialize a lazy value on the same parent object + * This causes scala to also try to grab a lock on the parent object. + * c) If thread 1 waits for thread 2 to join, a deadlock occurs. + * This wrapper will only grab a lock on the wrapper itself, and not the parent object. + * + * @param initialize The block of code to initialize the lazy value. + * @tparam T type of the lazy value. + */ +private[spark] class LazyTry[T](initialize: => T) extends Serializable { + private lazy val tryT: Try[T] = Utils.doTryWithCallerStacktrace { initialize } + + /** + * Get the lazy value. If the initialization block threw an exception, it will be re-thrown here. + * The exception will be re-thrown with the current caller's stacktrace. + * An exception with stack trace from when the exception was first thrown can be accessed with + * ``` + * ex.getSuppressed.find { e => + * e.getMessage == org.apache.spark.util.Utils.TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE + * } + * ``` + */ + def get: T = Utils.getTryWithCallerStacktrace(tryT) +} + +private[spark] object LazyTry { + /** + * Create a new LazyTry instance. + * + * @param initialize The block of code to initialize the lazy value. + * @tparam T type of the lazy value. + * @return a new LazyTry instance. + */ + def apply[T](initialize: => T): LazyTry[T] = new LazyTry(initialize) +} diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala index d8392cd8043de..52213f36a2cd1 100644 --- a/core/src/main/scala/org/apache/spark/util/Utils.scala +++ b/core/src/main/scala/org/apache/spark/util/Utils.scala @@ -1351,6 +1351,86 @@ private[spark] object Utils } } + val TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE = + "Full stacktrace of original doTryWithCallerStacktrace caller" + + val TRY_WITH_CALLER_STACKTRACE_TRY_STACKTRACE = + "Stacktrace under doTryWithCallerStacktrace" + + /** + * Use Try with stacktrace substitution for the caller retrieving the error. + * + * Normally in case of failure, the exception would have the stacktrace of the caller that + * originally called doTryWithCallerStacktrace. However, we want to replace the part above + * this function with the stacktrace of the caller who calls getTryWithCallerStacktrace. + * So here we save the part of the stacktrace below doTryWithCallerStacktrace, and + * getTryWithCallerStacktrace will stitch it with the new stack trace of the caller. + * The full original stack trace is kept in ex.getSuppressed. + * + * @param f Code block to be wrapped in Try + * @return Try with Success or Failure of the code block. Use with getTryWithCallerStacktrace. + */ + def doTryWithCallerStacktrace[T](f: => T): Try[T] = { + val t = Try { + f + } + t match { + case Failure(ex) => + // Note: we remove the common suffix instead of e.g. finding the call to this function, to + // account for recursive calls with multiple doTryWithCallerStacktrace on the stack trace. + val origStackTrace = ex.getStackTrace + val currentStackTrace = Thread.currentThread().getStackTrace + val commonSuffixLen = origStackTrace.reverse.zip(currentStackTrace.reverse).takeWhile { + case (exElem, currentElem) => exElem == currentElem + }.length + val belowEx = new Exception(TRY_WITH_CALLER_STACKTRACE_TRY_STACKTRACE) + belowEx.setStackTrace(origStackTrace.dropRight(commonSuffixLen)) + ex.addSuppressed(belowEx) + + // keep the full original stack trace in a suppressed exception. + val fullEx = new Exception(TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE) + fullEx.setStackTrace(origStackTrace) + ex.addSuppressed(fullEx) + case Success(_) => // nothing + } + t + } + + /** + * Retrieve the result of Try that was created by doTryWithCallerStacktrace. + * + * In case of failure, the resulting exception has a stack trace that combines the stack trace + * below the original doTryWithCallerStacktrace which triggered it, with the caller stack trace + * of the current caller of getTryWithCallerStacktrace. + * + * Full stack trace of the original doTryWithCallerStacktrace caller can be retrieved with + * ``` + * ex.getSuppressed.find { e => + * e.getMessage == Utils.TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE + * } + * ``` + * + * + * @param t Try from doTryWithCallerStacktrace + * @return Result of the Try or rethrows the failure exception with modified stacktrace. + */ + def getTryWithCallerStacktrace[T](t: Try[T]): T = t match { + case Failure(ex) => + val belowStacktrace = ex.getSuppressed.find { e => + // added in doTryWithCallerStacktrace + e.getMessage == TRY_WITH_CALLER_STACKTRACE_TRY_STACKTRACE + }.getOrElse { + // If we don't have the expected stacktrace information, just rethrow + throw ex + }.getStackTrace + // We are modifying and throwing the original exception. It would be better if we could + // return a copy, but we can't easily clone it and preserve. If this is accessed from + // multiple threads that then look at the stack trace, this could break. + ex.setStackTrace(belowStacktrace ++ Thread.currentThread().getStackTrace.drop(1)) + throw ex + case Success(s) => s + } + // A regular expression to match classes of the internal Spark API's // that we want to skip when finding the call site of a method. private val SPARK_CORE_CLASS_REGEX = diff --git a/core/src/test/scala/org/apache/spark/util/LazyTrySuite.scala b/core/src/test/scala/org/apache/spark/util/LazyTrySuite.scala new file mode 100644 index 0000000000000..79c07f8fbfead --- /dev/null +++ b/core/src/test/scala/org/apache/spark/util/LazyTrySuite.scala @@ -0,0 +1,151 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.util + +import org.apache.spark.SparkFunSuite + +class LazyTrySuite extends SparkFunSuite{ + test("LazyTry should initialize only once") { + var count = 0 + val lazyVal = LazyTry { + count += 1 + count + } + assert(count == 0) + assert(lazyVal.get == 1) + assert(count == 1) + assert(lazyVal.get == 1) + assert(count == 1) + } + + test("LazyTry should re-throw exceptions") { + val lazyVal = LazyTry { + throw new RuntimeException("test") + } + intercept[RuntimeException] { + lazyVal.get + } + intercept[RuntimeException] { + lazyVal.get + } + } + + test("LazyTry should re-throw exceptions with current caller stack-trace") { + val fileName = Thread.currentThread().getStackTrace()(1).getFileName + val lineNo = Thread.currentThread().getStackTrace()(1).getLineNumber + val lazyVal = LazyTry { + throw new RuntimeException("test") + } + + val e1 = intercept[RuntimeException] { + lazyVal.get // lineNo + 6 + } + assert(e1.getStackTrace + .exists(elem => elem.getFileName == fileName && elem.getLineNumber == lineNo + 6)) + + val e2 = intercept[RuntimeException] { + lazyVal.get // lineNo + 12 + } + assert(e2.getStackTrace + .exists(elem => elem.getFileName == fileName && elem.getLineNumber == lineNo + 12)) + } + + test("LazyTry does not lock containing object") { + class LazyContainer() { + @volatile var aSet = 0 + + val a: LazyTry[Int] = LazyTry { + aSet = 1 + aSet + } + + val b: LazyTry[Int] = LazyTry { + val t = new Thread(new Runnable { + override def run(): Unit = { + assert(a.get == 1) + } + }) + t.start() + t.join() + aSet + } + } + val container = new LazyContainer() + // Nothing is lazy initialized yet + assert(container.aSet == 0) + // This will not deadlock, thread t will initialize a, and update aSet + assert(container.b.get == 1) + assert(container.aSet == 1) + } + + // Scala lazy val tests are added to test for potential changes in the semantics of scala lazy val + + test("Scala lazy val initializing multiple times on error") { + class LazyValError() { + var counter = 0 + lazy val a = { + counter += 1 + throw new RuntimeException("test") + } + } + val lazyValError = new LazyValError() + intercept[RuntimeException] { + lazyValError.a + } + assert(lazyValError.counter == 1) + intercept[RuntimeException] { + lazyValError.a + } + assert(lazyValError.counter == 2) + } + + test("Scala lazy val locking containing object and deadlocking") { + // Note: this will change in scala 3, with different lazy vals not deadlocking with each other. + // https://docs.scala-lang.org/scala3/reference/changed-features/lazy-vals-init.html + class LazyValContainer() { + @volatile var aSet = 0 + @volatile var t: Thread = _ + + lazy val a = { + aSet = 1 + aSet + } + + lazy val b = { + t = new Thread(new Runnable { + override def run(): Unit = { + assert(a == 1) + } + }) + t.start() + t.join(1000) + aSet + } + } + val container = new LazyValContainer() + // Nothing is lazy initialized yet + assert(container.aSet == 0) + // This will deadlock, because b will take monitor on LazyValContainer, and then thread t + // will wait on that monitor, not able to initialize a. + // b will therefore see aSet == 0. + assert(container.b == 0) + // However, after b finishes initializing, the monitor will be released, and then thread t + // will finish initializing a, and set aSet to 1. + container.t.join() + assert(container.aSet == 1) + } +} diff --git a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala index 4fe6fcf17f49f..a694e08def89c 100644 --- a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala @@ -28,7 +28,7 @@ import java.util.concurrent.TimeUnit import java.util.zip.GZIPOutputStream import scala.collection.mutable.ListBuffer -import scala.util.Random +import scala.util.{Random, Try} import com.google.common.io.Files import org.apache.commons.io.IOUtils @@ -1523,6 +1523,116 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties { conf.set(SERIALIZER, "org.apache.spark.serializer.JavaSerializer") assert(Utils.isPushBasedShuffleEnabled(conf, isDriver = true) === false) } + + + private def throwException(): String = { + throw new Exception("test") + } + + private def callDoTry(): Try[String] = { + Utils.doTryWithCallerStacktrace { + throwException() + } + } + + private def callGetTry(t: Try[String]): String = { + Utils.getTryWithCallerStacktrace(t) + } + + private def callGetTryAgain(t: Try[String]): String = { + Utils.getTryWithCallerStacktrace(t) + } + + test("doTryWithCallerStacktrace and getTryWithCallerStacktrace") { + val t = callDoTry() + + val e1 = intercept[Exception] { + callGetTry(t) + } + // Uncomment for manual inspection + // e1.printStackTrace() + // Example: + // java.lang.Exception: test + // at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1640) + // at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1645) + // at scala.util.Try$.apply(Try.scala:213) + // at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1586) + // at org.apache.spark.util.Utils$.getTryWithCallerStacktrace(Utils.scala:1639) + // at org.apache.spark.util.UtilsSuite.callGetTry(UtilsSuite.scala:1650) + // at org.apache.spark.util.UtilsSuite.$anonfun$new$165(UtilsSuite.scala:1661) + // <- callGetTry is seen as calling getTryWithCallerStacktrace + + val st1 = e1.getStackTrace + // throwException should be on the stack trace + assert(st1.exists(_.getMethodName == "throwException")) + // callDoTry shouldn't be on the stack trace, but callGetTry should be. + assert(!st1.exists(_.getMethodName == "callDoTry")) + assert(st1.exists(_.getMethodName == "callGetTry")) + + // The original stack trace with callDoTry should be in the suppressed exceptions. + // Example: + // scalastyle:off line.size.limit + // Suppressed: java.lang.Exception: Full stacktrace of original doTryWithCallerStacktrace caller + // at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1640) + // at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1645) + // at scala.util.Try$.apply(Try.scala:213) + // at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1586) + // at org.apache.spark.util.UtilsSuite.callDoTry(UtilsSuite.scala:1645) + // at org.apache.spark.util.UtilsSuite.$anonfun$new$165(UtilsSuite.scala:1658) + // ... 56 more + // scalastyle:on line.size.limit + val origSt = e1.getSuppressed.find( + _.getMessage == Utils.TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE) + assert(origSt.isDefined) + assert(origSt.get.getStackTrace.exists(_.getMethodName == "throwException")) + assert(origSt.get.getStackTrace.exists(_.getMethodName == "callDoTry")) + + // The stack trace under Try should be in the suppressed exceptions. + // Example: + // Suppressed: java.lang.Exception: Stacktrace under doTryWithCallerStacktrace + // at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala: 1640) + // at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala: 1645) + // at scala.util.Try$.apply(Try.scala: 213) + // at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala: 1586) + val trySt = e1.getSuppressed.find( + _.getMessage == Utils.TRY_WITH_CALLER_STACKTRACE_TRY_STACKTRACE) + assert(trySt.isDefined) + // calls under callDoTry should be present. + assert(trySt.get.getStackTrace.exists(_.getMethodName == "throwException")) + // callDoTry should be removed. + assert(!trySt.get.getStackTrace.exists(_.getMethodName == "callDoTry")) + + val e2 = intercept[Exception] { + callGetTryAgain(t) + } + // Uncomment for manual inspection + // e2.printStackTrace() + // Example: + // java.lang.Exception: test + // at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1640) + // at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1645) + // at scala.util.Try$.apply(Try.scala:213) + // at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1586) + // at org.apache.spark.util.Utils$.getTryWithCallerStacktrace(Utils.scala:1639) + // at org.apache.spark.util.UtilsSuite.callGetTryAgain(UtilsSuite.scala:1654) + // at org.apache.spark.util.UtilsSuite.$anonfun$new$165(UtilsSuite.scala:1711) + // <- callGetTryAgain is seen as calling getTryWithCallerStacktrace + + val st2 = e2.getStackTrace + // throwException should be on the stack trace + assert(st2.exists(_.getMethodName == "throwException")) + // callDoTry shouldn't be on the stack trace, but callGetTryAgain should be. + assert(!st2.exists(_.getMethodName == "callDoTry")) + assert(st2.exists(_.getMethodName == "callGetTryAgain")) + // callGetTry that we called before shouldn't be on the stack trace. + assert(!st2.exists(_.getMethodName == "callGetTry")) + + // Unfortunately, this utility is not able to clone the exception, but modifies it in place, + // so now e1 is also pointing to "callGetTryAgain" instead of "callGetTry". + val st1Again = e1.getStackTrace + assert(st1Again.exists(_.getMethodName == "callGetTryAgain")) + assert(!st1Again.exists(_.getMethodName == "callGetTry")) + } } private class SimpleExtension diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala index dc918e51d0550..2115e21f81d71 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala @@ -60,6 +60,11 @@ case class CollectMetricsExec( override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override def resetMetrics(): Unit = { + accumulator.reset() + super.resetMetrics() + } + override protected def doExecute(): RDD[InternalRow] = { val collector = accumulator collector.reset() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala index 7bc770a0c9e33..fb3ec3ad41812 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.execution.datasources.WriteFilesSpec import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.NextIterator +import org.apache.spark.util.{LazyTry, NextIterator} import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} object SparkPlan { @@ -182,6 +182,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ /** Specifies sort order for each partition requirements on the input data for this operator. */ def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq.fill(children.size)(Nil) + @transient + private val executeRDD = LazyTry { + doExecute() + } + /** * Returns the result of this query as an RDD[InternalRow] by delegating to `doExecute` after * preparations. @@ -192,7 +197,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ if (isCanonicalizedPlan) { throw SparkException.internalError("A canonicalized plan is not supposed to be executed.") } - doExecute() + executeRDD.get + } + + private val executeBroadcastBcast = LazyTry { + doExecuteBroadcast() } /** @@ -205,7 +214,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ if (isCanonicalizedPlan) { throw SparkException.internalError("A canonicalized plan is not supposed to be executed.") } - doExecuteBroadcast() + executeBroadcastBcast.get.asInstanceOf[broadcast.Broadcast[T]] + } + + private val executeColumnarRDD = LazyTry { + doExecuteColumnar() } /** @@ -219,7 +232,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ if (isCanonicalizedPlan) { throw SparkException.internalError("A canonicalized plan is not supposed to be executed.") } - doExecuteColumnar() + executeColumnarRDD.get } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala index cfcfd282e5480..cbd60804b27e8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala @@ -99,48 +99,6 @@ case class InMemoryTableScanExec( relation.cacheBuilder.serializer.supportsColumnarOutput(relation.schema) } - private lazy val columnarInputRDD: RDD[ColumnarBatch] = { - val numOutputRows = longMetric("numOutputRows") - val buffers = filteredCachedBatches() - relation.cacheBuilder.serializer.convertCachedBatchToColumnarBatch( - buffers, - relation.output, - attributes, - conf).map { cb => - numOutputRows += cb.numRows() - cb - } - } - - private lazy val inputRDD: RDD[InternalRow] = { - if (enableAccumulatorsForTest) { - readPartitions.setValue(0) - readBatches.setValue(0) - } - - val numOutputRows = longMetric("numOutputRows") - // Using these variables here to avoid serialization of entire objects (if referenced - // directly) within the map Partitions closure. - val relOutput = relation.output - val serializer = relation.cacheBuilder.serializer - - // update SQL metrics - val withMetrics = - filteredCachedBatches().mapPartitionsInternal { iter => - if (enableAccumulatorsForTest && iter.hasNext) { - readPartitions.add(1) - } - iter.map { batch => - if (enableAccumulatorsForTest) { - readBatches.add(1) - } - numOutputRows += batch.numRows - batch - } - } - serializer.convertCachedBatchToInternalRow(withMetrics, relOutput, attributes, conf) - } - override def output: Seq[Attribute] = attributes private def cachedPlan = relation.cachedPlan match { @@ -191,11 +149,47 @@ case class InMemoryTableScanExec( } protected override def doExecute(): RDD[InternalRow] = { - inputRDD + // Resulting RDD is cached and reused by SparkPlan.executeRDD + if (enableAccumulatorsForTest) { + readPartitions.setValue(0) + readBatches.setValue(0) + } + + val numOutputRows = longMetric("numOutputRows") + // Using these variables here to avoid serialization of entire objects (if referenced + // directly) within the map Partitions closure. + val relOutput = relation.output + val serializer = relation.cacheBuilder.serializer + + // update SQL metrics + val withMetrics = + filteredCachedBatches().mapPartitionsInternal { iter => + if (enableAccumulatorsForTest && iter.hasNext) { + readPartitions.add(1) + } + iter.map { batch => + if (enableAccumulatorsForTest) { + readBatches.add(1) + } + numOutputRows += batch.numRows + batch + } + } + serializer.convertCachedBatchToInternalRow(withMetrics, relOutput, attributes, conf) } protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { - columnarInputRDD + // Resulting RDD is cached and reused by SparkPlan.executeColumnarRDD + val numOutputRows = longMetric("numOutputRows") + val buffers = filteredCachedBatches() + relation.cacheBuilder.serializer.convertCachedBatchToColumnarBatch( + buffers, + relation.output, + attributes, + conf).map { cb => + numOutputRows += cb.numRows() + cb + } } override def isMaterialized: Boolean = relation.cacheBuilder.isCachedColumnBuffersLoaded diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index 90f00a5035e15..ae11229cd516e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -249,17 +249,10 @@ case class ShuffleExchangeExec( dep } - /** - * Caches the created ShuffleRowRDD so we can reuse that. - */ - private var cachedShuffleRDD: ShuffledRowRDD = null - protected override def doExecute(): RDD[InternalRow] = { - // Returns the same ShuffleRowRDD if this plan is used by multiple plans. - if (cachedShuffleRDD == null) { - cachedShuffleRDD = new ShuffledRowRDD(shuffleDependency, readMetrics) - } - cachedShuffleRDD + // The ShuffleRowRDD will be cached in SparkPlan.executeRDD and reused if this plan is used by + // multiple plans. + new ShuffledRowRDD(shuffleDependency, readMetrics) } override protected def withNewChildInternal(newChild: SparkPlan): ShuffleExchangeExec = From 44ec70f5103fc5674497373ac5c23e8145ae5660 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Mon, 23 Sep 2024 18:28:19 +0800 Subject: [PATCH 106/189] [SPARK-49626][PYTHON][CONNECT] Support horizontal and vertical bar plots ### What changes were proposed in this pull request? Support horizontal and vertical bar plots with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. ```python >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] >>> columns = ["category", "int_val", "float_val"] >>> sdf = spark.createDataFrame(data, columns) >>> sdf.show() +--------+-------+---------+ |category|int_val|float_val| +--------+-------+---------+ | A| 10| 1.5| | B| 30| 2.5| | C| 20| 3.5| +--------+-------+---------+ >>> f = sdf.plot(kind="bar", x="category", y=["int_val", "float_val"]) >>> f.show() # see below >>> g = sdf.plot.barh(x=["int_val", "float_val"], y="category") >>> g.show() # see below ``` `f.show()`: ![newplot (4)](https://github.com/user-attachments/assets/0df9ee86-fb48-4796-b6c3-aaf2879217aa) `g.show()`: ![newplot (3)](https://github.com/user-attachments/assets/f39b01c3-66e6-464b-b2e8-badebb39bc67) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48100 from xinrong-meng/plot_bar. Authored-by: Xinrong Meng Signed-off-by: Xinrong Meng --- python/pyspark/sql/plot/core.py | 79 +++++++++++++++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 44 +++++++++-- 2 files changed, 117 insertions(+), 6 deletions(-) diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index 392ef73b38845..ed22d02370ca6 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -75,6 +75,8 @@ def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": class PySparkPlotAccessor: plot_data_map = { + "bar": PySparkTopNPlotBase().get_top_n, + "barh": PySparkTopNPlotBase().get_top_n, "line": PySparkSampledPlotBase().get_sampled, } _backends = {} # type: ignore[var-annotated] @@ -133,3 +135,80 @@ def line(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": >>> df.plot.line(x="category", y=["int_val", "float_val"]) # doctest: +SKIP """ return self(kind="line", x=x, y=y, **kwargs) + + def bar(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": + """ + Vertical bar plot. + + A bar plot is a plot that presents categorical data with rectangular bars with lengths + proportional to the values that they represent. A bar plot shows comparisons among + discrete categories. One axis of the plot shows the specific categories being compared, + and the other axis represents a measured value. + + Parameters + ---------- + x : str + Name of column to use for the horizontal axis. + y : str or list of str + Name(s) of the column(s) to use for the vertical axis. + Multiple columns can be plotted. + **kwargs : optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.bar(x="category", y="int_val") # doctest: +SKIP + >>> df.plot.bar(x="category", y=["int_val", "float_val"]) # doctest: +SKIP + """ + return self(kind="bar", x=x, y=y, **kwargs) + + def barh(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": + """ + Make a horizontal bar plot. + + A horizontal bar plot is a plot that presents quantitative data with + rectangular bars with lengths proportional to the values that they + represent. A bar plot shows comparisons among discrete categories. One + axis of the plot shows the specific categories being compared, and the + other axis represents a measured value. + + Parameters + ---------- + x : str or list of str + Name(s) of the column(s) to use for the horizontal axis. + Multiple columns can be plotted. + y : str or list of str + Name(s) of the column(s) to use for the vertical axis. + Multiple columns can be plotted. + **kwargs : optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Notes + ----- + In Plotly and Matplotlib, the interpretation of `x` and `y` for `barh` plots differs. + In Plotly, `x` refers to the values and `y` refers to the categories. + In Matplotlib, `x` refers to the categories and `y` refers to the values. + Ensure correct axis labeling based on the backend used. + + Examples + -------- + >>> data = [("A", 10, 1.5), ("B", 30, 2.5), ("C", 20, 3.5)] + >>> columns = ["category", "int_val", "float_val"] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.barh(x="int_val", y="category") # doctest: +SKIP + >>> df.plot.barh( + ... x=["int_val", "float_val"], y="category" + ... ) # doctest: +SKIP + """ + return self(kind="barh", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py index 72a3ed267d192..1c52c93a23d3a 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -28,9 +28,16 @@ def sdf(self): columns = ["category", "int_val", "float_val"] return self.spark.createDataFrame(data, columns) - def _check_fig_data(self, fig_data, expected_x, expected_y, expected_name=""): - self.assertEqual(fig_data["mode"], "lines") - self.assertEqual(fig_data["type"], "scatter") + def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name=""): + if kind == "line": + self.assertEqual(fig_data["mode"], "lines") + self.assertEqual(fig_data["type"], "scatter") + elif kind == "bar": + self.assertEqual(fig_data["type"], "bar") + elif kind == "barh": + self.assertEqual(fig_data["type"], "bar") + self.assertEqual(fig_data["orientation"], "h") + self.assertEqual(fig_data["xaxis"], "x") self.assertEqual(list(fig_data["x"]), expected_x) self.assertEqual(fig_data["yaxis"], "y") @@ -40,12 +47,37 @@ def _check_fig_data(self, fig_data, expected_x, expected_y, expected_name=""): def test_line_plot(self): # single column as vertical axis fig = self.sdf.plot(kind="line", x="category", y="int_val") - self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20]) + self._check_fig_data("line", fig["data"][0], ["A", "B", "C"], [10, 30, 20]) # multiple columns as vertical axis fig = self.sdf.plot.line(x="category", y=["int_val", "float_val"]) - self._check_fig_data(fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") - self._check_fig_data(fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + self._check_fig_data("line", fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") + self._check_fig_data("line", fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + + def test_bar_plot(self): + # single column as vertical axis + fig = self.sdf.plot(kind="bar", x="category", y="int_val") + self._check_fig_data("bar", fig["data"][0], ["A", "B", "C"], [10, 30, 20]) + + # multiple columns as vertical axis + fig = self.sdf.plot.bar(x="category", y=["int_val", "float_val"]) + self._check_fig_data("bar", fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") + self._check_fig_data("bar", fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + + def test_barh_plot(self): + # single column as vertical axis + fig = self.sdf.plot(kind="barh", x="category", y="int_val") + self._check_fig_data("barh", fig["data"][0], ["A", "B", "C"], [10, 30, 20]) + + # multiple columns as vertical axis + fig = self.sdf.plot.barh(x="category", y=["int_val", "float_val"]) + self._check_fig_data("barh", fig["data"][0], ["A", "B", "C"], [10, 30, 20], "int_val") + self._check_fig_data("barh", fig["data"][1], ["A", "B", "C"], [1.5, 2.5, 3.5], "float_val") + + # multiple columns as horizontal axis + fig = self.sdf.plot.barh(x=["int_val", "float_val"], y="category") + self._check_fig_data("barh", fig["data"][0], [10, 30, 20], ["A", "B", "C"], "int_val") + self._check_fig_data("barh", fig["data"][1], [1.5, 2.5, 3.5], ["A", "B", "C"], "float_val") class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): From e1637e3fbe0a7ee6492cfc909ef13fc1fe0534d1 Mon Sep 17 00:00:00 2001 From: Kent Yao Date: Mon, 23 Sep 2024 19:51:21 +0800 Subject: [PATCH 107/189] [SPARK-48712][SQL][FOLLOWUP] Check whether input is valid utf-8 string or not before entering fast path ### What changes were proposed in this pull request? Check whether input is valid utf-8 string or not before entering fast path ### Why are the changes needed? Avoid behavior change on a corner case where users provide invalid UTF-8 strings for UTF-8 encoding ### Does this PR introduce _any_ user-facing change? no, this is a followup to avoid potential breaking change ### How was this patch tested? existing tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48203 from yaooqinn/SPARK-48712. Authored-by: Kent Yao Signed-off-by: Kent Yao --- .../expressions/stringExpressions.scala | 5 ++--- .../expressions/StringExpressionsSuite.scala | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index da6d786efb4e3..786c3968be0fe 100755 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -3039,10 +3039,9 @@ object Encode { legacyCharsets: Boolean, legacyErrorAction: Boolean): Array[Byte] = { val toCharset = charset.toString - if (input.numBytes == 0 || "UTF-8".equalsIgnoreCase(toCharset)) { - return input.getBytes - } + if ("UTF-8".equalsIgnoreCase(toCharset) && input.isValid) return input.getBytes val encoder = CharsetProvider.newEncoder(toCharset, legacyCharsets, legacyErrorAction) + if (input.numBytes == 0) return input.getBytes try { val bb = encoder.encode(CharBuffer.wrap(input.toString)) JavaUtils.bufferToArray(bb) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala index 29b878230472d..9b454ba764f92 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala @@ -26,9 +26,12 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions.Cast._ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection +import org.apache.spark.sql.catalyst.util.CharsetProvider +import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLId import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.types.StringTypeAnyCollation import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { @@ -2076,4 +2079,22 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { ) ) } + + test("SPARK-48712: Check whether input is valid utf-8 string or not before entering fast path") { + val str = UTF8String.fromBytes(Array[Byte](-1, -2, -3, -4)) + assert(!str.isValid, "please use a string that is not valid UTF-8 for testing") + val expected = Array[Byte](-17, -65, -67, -17, -65, -67, -17, -65, -67, -17, -65, -67) + val bytes = Encode.encode(str, UTF8String.fromString("UTF-8"), false, false) + assert(bytes === expected) + checkEvaluation(Encode(Literal(str), Literal("UTF-8")), expected) + checkEvaluation(Encode(Literal(UTF8String.EMPTY_UTF8), Literal("UTF-8")), Array.emptyByteArray) + checkErrorInExpression[SparkIllegalArgumentException]( + Encode(Literal(UTF8String.EMPTY_UTF8), Literal("UTF-12345")), + condition = "INVALID_PARAMETER_VALUE.CHARSET", + parameters = Map( + "charset" -> "UTF-12345", + "functionName" -> toSQLId("encode"), + "parameter" -> toSQLId("charset"), + "charsets" -> CharsetProvider.VALID_CHARSETS.mkString(", "))) + } } From fec1562b0ea03ff42d2468ea8ff7cbbc569336d8 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Mon, 23 Sep 2024 20:03:14 +0800 Subject: [PATCH 108/189] [SPARK-49755][CONNECT] Remove special casing for avro functions in Connect ### What changes were proposed in this pull request? apply the built-in registered functions ### Why are the changes needed? code simplification ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? updated tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48209 from zhengruifeng/connect_avro. Authored-by: Ruifeng Zheng Signed-off-by: yangjie01 --- .../expressions/toFromAvroSqlFunctions.scala | 3 ++ .../from_avro_with_options.explain | 2 +- .../from_avro_without_options.explain | 2 +- .../to_avro_with_schema.explain | 2 +- .../to_avro_without_schema.explain | 2 +- sql/connect/server/pom.xml | 2 +- .../connect/planner/SparkConnectPlanner.scala | 47 +------------------ 7 files changed, 9 insertions(+), 51 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromAvroSqlFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromAvroSqlFunctions.scala index 58bddafac0882..457f469e0f687 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromAvroSqlFunctions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromAvroSqlFunctions.scala @@ -61,6 +61,9 @@ case class FromAvro(child: Expression, jsonFormatSchema: Expression, options: Ex override def second: Expression = jsonFormatSchema override def third: Expression = options + def this(child: Expression, jsonFormatSchema: Expression) = + this(child, jsonFormatSchema, Literal.create(null)) + override def withNewChildrenInternal( newFirst: Expression, newSecond: Expression, newThird: Expression): Expression = { copy(child = newFirst, jsonFormatSchema = newSecond, options = newThird) diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_with_options.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_with_options.explain index 1ef91ef8c36ac..f08c804d3b88a 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_with_options.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_with_options.explain @@ -1,2 +1,2 @@ -Project [from_avro(bytes#0, {"type": "int", "name": "id"}, (mode,FAILFAST), (compression,zstandard)) AS from_avro(bytes)#0] +Project [from_avro(bytes#0, {"type": "int", "name": "id"}, (mode,FAILFAST), (compression,zstandard)) AS from_avro(bytes, {"type": "int", "name": "id"}, map(mode, FAILFAST, compression, zstandard))#0] +- LocalRelation , [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_without_options.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_without_options.explain index 8fca0b5341694..6fe4a8babc689 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_without_options.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/from_avro_without_options.explain @@ -1,2 +1,2 @@ -Project [from_avro(bytes#0, {"type": "string", "name": "name"}) AS from_avro(bytes)#0] +Project [from_avro(bytes#0, {"type": "string", "name": "name"}) AS from_avro(bytes, {"type": "string", "name": "name"}, NULL)#0] +- LocalRelation , [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_with_schema.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_with_schema.explain index cd2dc984e3ffa..8ba9248f844c7 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_with_schema.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_with_schema.explain @@ -1,2 +1,2 @@ -Project [to_avro(a#0, Some({"type": "int", "name": "id"})) AS to_avro(a)#0] +Project [to_avro(a#0, Some({"type": "int", "name": "id"})) AS to_avro(a, {"type": "int", "name": "id"})#0] +- LocalRelation , [id#0L, a#0, b#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_without_schema.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_without_schema.explain index a5371c70ac78a..b2947334945e3 100644 --- a/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_without_schema.explain +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/to_avro_without_schema.explain @@ -1,2 +1,2 @@ -Project [to_avro(id#0L, None) AS to_avro(id)#0] +Project [to_avro(id#0L, None) AS to_avro(id, NULL)#0] +- LocalRelation , [id#0L, a#0, b#0] diff --git a/sql/connect/server/pom.xml b/sql/connect/server/pom.xml index 3350c4261e9da..12e3ed9030437 100644 --- a/sql/connect/server/pom.xml +++ b/sql/connect/server/pom.xml @@ -105,7 +105,7 @@ org.apache.spark spark-avro_${scala.binary.version} ${project.version} - provided + test org.apache.spark diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala index 33c9edb1cd21a..231e54ff77d29 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala @@ -44,7 +44,6 @@ import org.apache.spark.internal.{Logging, LogKeys, MDC} import org.apache.spark.internal.LogKeys.{DATAFRAME_ID, SESSION_ID} import org.apache.spark.resource.{ExecutorResourceRequest, ResourceProfile, TaskResourceProfile, TaskResourceRequest} import org.apache.spark.sql.{Dataset, Encoders, ForeachWriter, Observation, RelationalGroupedDataset, Row, SparkSession} -import org.apache.spark.sql.avro.{AvroDataToCatalyst, CatalystDataToAvro} import org.apache.spark.sql.catalyst.{expressions, AliasIdentifier, FunctionIdentifier, QueryPlanningTracker} import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, GlobalTempView, LocalTempView, MultiAlias, NameParameterizedQuery, PosParameterizedQuery, UnresolvedAlias, UnresolvedAttribute, UnresolvedDataFrameStar, UnresolvedDeserializer, UnresolvedExtractValue, UnresolvedFunction, UnresolvedRegex, UnresolvedRelation, UnresolvedStar, UnresolvedTranspose} import org.apache.spark.sql.catalyst.encoders.{encoderFor, AgnosticEncoder, ExpressionEncoder, RowEncoder} @@ -1523,8 +1522,7 @@ class SparkConnectPlanner( case proto.Expression.ExprTypeCase.UNRESOLVED_ATTRIBUTE => transformUnresolvedAttribute(exp.getUnresolvedAttribute) case proto.Expression.ExprTypeCase.UNRESOLVED_FUNCTION => - transformUnregisteredFunction(exp.getUnresolvedFunction) - .getOrElse(transformUnresolvedFunction(exp.getUnresolvedFunction)) + transformUnresolvedFunction(exp.getUnresolvedFunction) case proto.Expression.ExprTypeCase.ALIAS => transformAlias(exp.getAlias) case proto.Expression.ExprTypeCase.EXPRESSION_STRING => transformExpressionString(exp.getExpressionString) @@ -1844,49 +1842,6 @@ class SparkConnectPlanner( UnresolvedNamedLambdaVariable(variable.getNamePartsList.asScala.toSeq) } - /** - * For some reason, not all functions are registered in 'FunctionRegistry'. For a unregistered - * function, we can still wrap it under the proto 'UnresolvedFunction', and then resolve it in - * this method. - */ - private def transformUnregisteredFunction( - fun: proto.Expression.UnresolvedFunction): Option[Expression] = { - fun.getFunctionName match { - // Avro-specific functions - case "from_avro" if Seq(2, 3).contains(fun.getArgumentsCount) => - val children = fun.getArgumentsList.asScala.map(transformExpression) - val jsonFormatSchema = extractString(children(1), "jsonFormatSchema") - var options = Map.empty[String, String] - if (fun.getArgumentsCount == 3) { - options = extractMapData(children(2), "Options") - } - Some(AvroDataToCatalyst(children.head, jsonFormatSchema, options)) - - case "to_avro" if Seq(1, 2).contains(fun.getArgumentsCount) => - val children = fun.getArgumentsList.asScala.map(transformExpression) - var jsonFormatSchema = Option.empty[String] - if (fun.getArgumentsCount == 2) { - jsonFormatSchema = Some(extractString(children(1), "jsonFormatSchema")) - } - Some(CatalystDataToAvro(children.head, jsonFormatSchema)) - - case _ => None - } - } - - private def extractString(expr: Expression, field: String): String = expr match { - case Literal(s, StringType) if s != null => s.toString - case other => throw InvalidPlanInput(s"$field should be a literal string, but got $other") - } - - @scala.annotation.tailrec - private def extractMapData(expr: Expression, field: String): Map[String, String] = expr match { - case map: CreateMap => ExprUtils.convertToMapData(map) - case UnresolvedFunction(Seq("map"), args, _, _, _, _, _) => - extractMapData(CreateMap(args), field) - case other => throw InvalidPlanInput(s"$field should be created by map, but got $other") - } - private def transformAlias(alias: proto.Expression.Alias): NamedExpression = { if (alias.getNameCount == 1) { val metadata = if (alias.hasMetadata() && alias.getMetadata.nonEmpty) { From 3b5c1d6baeb239c75c182513b3fad37d532d9f9f Mon Sep 17 00:00:00 2001 From: Nemanja Boric Date: Mon, 23 Sep 2024 11:22:07 -0400 Subject: [PATCH 109/189] [SPARK-49747][CONNECT] Migrate connect/ files to structured logging ### What changes were proposed in this pull request? We are moving one missing piece in SparkConnect to MDC-based logging. ### Why are the changes needed? As part of the greater migration to structured logging. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Compilation/existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48195 from nemanja-boric-databricks/mdc-connect. Authored-by: Nemanja Boric Signed-off-by: Herman van Hovell --- .../spark/sql/connect/execution/ExecuteThreadRunner.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala index fe43edb5c6218..e75654e2c384f 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala @@ -27,7 +27,7 @@ import org.apache.commons.lang3.StringUtils import org.apache.spark.SparkSQLException import org.apache.spark.connect.proto -import org.apache.spark.internal.Logging +import org.apache.spark.internal.{Logging, LogKeys, MDC} import org.apache.spark.sql.connect.common.ProtoUtils import org.apache.spark.sql.connect.planner.SparkConnectPlanner import org.apache.spark.sql.connect.service.{ExecuteHolder, ExecuteSessionTag, SparkConnectService} @@ -113,7 +113,7 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends } catch { // Need to catch throwable instead of NonFatal, because e.g. InterruptedException is fatal. case e: Throwable => - logDebug(s"Exception in execute: $e") + logDebug(log"Exception in execute: ${MDC(LogKeys.EXCEPTION, e)}") // Always cancel all remaining execution after error. executeHolder.sessionHolder.session.sparkContext.cancelJobsWithTag( executeHolder.jobTag, @@ -298,7 +298,7 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends ProtoUtils.abbreviate(request, maxLevel = 8).toString) } catch { case NonFatal(e) => - logWarning("Fail to extract debug information", e) + logWarning(log"Fail to extract debug information: ${MDC(LogKeys.EXCEPTION, e)}") "UNKNOWN" } } From 1086256a81f16127563cdf9a6d0b7ef1e413f17a Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Mon, 23 Sep 2024 19:10:44 -0400 Subject: [PATCH 110/189] [SPARK-49415][CONNECT][SQL] Move SQLImplicits to sql/api ### What changes were proposed in this pull request? This PR largely moves SQLImplicits and DatasetHolder to sql/api. ### Why are the changes needed? We are creating a unified Scala interface for Classic and Connect. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48151 from hvanhovell/SPARK-49415. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/DatasetHolder.scala | 41 --- .../org/apache/spark/sql/SQLImplicits.scala | 283 +---------------- .../org/apache/spark/sql/SparkSession.scala | 15 +- project/MimaExcludes.scala | 12 + .../org/apache/spark/sql/DatasetHolder.scala | 11 +- .../apache/spark/sql/api/SQLImplicits.scala | 300 ++++++++++++++++++ .../apache/spark/sql/api/SparkSession.scala | 13 + .../org/apache/spark/sql/SQLContext.scala | 4 +- .../org/apache/spark/sql/SQLImplicits.scala | 248 +-------------- .../org/apache/spark/sql/SparkSession.scala | 15 +- .../sql/expressions/scalalang/typed.scala | 5 - .../apache/spark/sql/test/SQLTestData.scala | 2 +- .../apache/spark/sql/test/SQLTestUtils.scala | 2 +- 13 files changed, 348 insertions(+), 603 deletions(-) delete mode 100644 connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DatasetHolder.scala rename sql/{core => api}/src/main/scala/org/apache/spark/sql/DatasetHolder.scala (79%) create mode 100644 sql/api/src/main/scala/org/apache/spark/sql/api/SQLImplicits.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DatasetHolder.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DatasetHolder.scala deleted file mode 100644 index 66f591bf1fb99..0000000000000 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/DatasetHolder.scala +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.sql - -/** - * A container for a [[Dataset]], used for implicit conversions in Scala. - * - * To use this, import implicit conversions in SQL: - * {{{ - * val spark: SparkSession = ... - * import spark.implicits._ - * }}} - * - * @since 3.4.0 - */ -case class DatasetHolder[T] private[sql] (private val ds: Dataset[T]) { - - // This is declared with parentheses to prevent the Scala compiler from treating - // `rdd.toDS("1")` as invoking this toDS and then apply on the returned Dataset. - def toDS(): Dataset[T] = ds - - // This is declared with parentheses to prevent the Scala compiler from treating - // `rdd.toDF("1")` as invoking this toDF and then apply on the returned DataFrame. - def toDF(): DataFrame = ds.toDF() - - def toDF(colNames: String*): DataFrame = ds.toDF(colNames: _*) -} diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SQLImplicits.scala index 7799d395d5c6a..4690253da808b 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SQLImplicits.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SQLImplicits.scala @@ -16,283 +16,8 @@ */ package org.apache.spark.sql -import scala.collection.Map -import scala.language.implicitConversions -import scala.reflect.classTag -import scala.reflect.runtime.universe.TypeTag - -import org.apache.spark.sql.catalyst.ScalaReflection -import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, AgnosticEncoders} -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders._ - -/** - * A collection of implicit methods for converting names and Symbols into [[Column]]s, and for - * converting common Scala objects into [[Dataset]]s. - * - * @since 3.4.0 - */ -abstract class SQLImplicits private[sql] (session: SparkSession) extends LowPrioritySQLImplicits { - - /** - * Converts $"col name" into a [[Column]]. - * - * @since 3.4.0 - */ - implicit class StringToColumn(val sc: StringContext) { - def $(args: Any*): ColumnName = { - new ColumnName(sc.s(args: _*)) - } - } - - /** - * An implicit conversion that turns a Scala `Symbol` into a [[Column]]. - * @since 3.4.0 - */ - implicit def symbolToColumn(s: Symbol): ColumnName = new ColumnName(s.name) - - /** @since 3.4.0 */ - implicit val newIntEncoder: Encoder[Int] = PrimitiveIntEncoder - - /** @since 3.4.0 */ - implicit val newLongEncoder: Encoder[Long] = PrimitiveLongEncoder - - /** @since 3.4.0 */ - implicit val newDoubleEncoder: Encoder[Double] = PrimitiveDoubleEncoder - - /** @since 3.4.0 */ - implicit val newFloatEncoder: Encoder[Float] = PrimitiveFloatEncoder - - /** @since 3.4.0 */ - implicit val newByteEncoder: Encoder[Byte] = PrimitiveByteEncoder - - /** @since 3.4.0 */ - implicit val newShortEncoder: Encoder[Short] = PrimitiveShortEncoder - - /** @since 3.4.0 */ - implicit val newBooleanEncoder: Encoder[Boolean] = PrimitiveBooleanEncoder - - /** @since 3.4.0 */ - implicit val newStringEncoder: Encoder[String] = StringEncoder - - /** @since 3.4.0 */ - implicit val newJavaDecimalEncoder: Encoder[java.math.BigDecimal] = - AgnosticEncoders.DEFAULT_JAVA_DECIMAL_ENCODER - - /** @since 3.4.0 */ - implicit val newScalaDecimalEncoder: Encoder[scala.math.BigDecimal] = - AgnosticEncoders.DEFAULT_SCALA_DECIMAL_ENCODER - - /** @since 3.4.0 */ - implicit val newDateEncoder: Encoder[java.sql.Date] = AgnosticEncoders.STRICT_DATE_ENCODER - - /** @since 3.4.0 */ - implicit val newLocalDateEncoder: Encoder[java.time.LocalDate] = - AgnosticEncoders.STRICT_LOCAL_DATE_ENCODER - - /** @since 3.4.0 */ - implicit val newLocalDateTimeEncoder: Encoder[java.time.LocalDateTime] = - AgnosticEncoders.LocalDateTimeEncoder - - /** @since 3.4.0 */ - implicit val newTimeStampEncoder: Encoder[java.sql.Timestamp] = - AgnosticEncoders.STRICT_TIMESTAMP_ENCODER - - /** @since 3.4.0 */ - implicit val newInstantEncoder: Encoder[java.time.Instant] = - AgnosticEncoders.STRICT_INSTANT_ENCODER - - /** @since 3.4.0 */ - implicit val newDurationEncoder: Encoder[java.time.Duration] = DayTimeIntervalEncoder - - /** @since 3.4.0 */ - implicit val newPeriodEncoder: Encoder[java.time.Period] = YearMonthIntervalEncoder - - /** @since 3.4.0 */ - implicit def newJavaEnumEncoder[A <: java.lang.Enum[_]: TypeTag]: Encoder[A] = { - ScalaReflection.encoderFor[A] - } - - // Boxed primitives - - /** @since 3.4.0 */ - implicit val newBoxedIntEncoder: Encoder[java.lang.Integer] = BoxedIntEncoder - - /** @since 3.4.0 */ - implicit val newBoxedLongEncoder: Encoder[java.lang.Long] = BoxedLongEncoder - - /** @since 3.4.0 */ - implicit val newBoxedDoubleEncoder: Encoder[java.lang.Double] = BoxedDoubleEncoder - - /** @since 3.4.0 */ - implicit val newBoxedFloatEncoder: Encoder[java.lang.Float] = BoxedFloatEncoder - - /** @since 3.4.0 */ - implicit val newBoxedByteEncoder: Encoder[java.lang.Byte] = BoxedByteEncoder - - /** @since 3.4.0 */ - implicit val newBoxedShortEncoder: Encoder[java.lang.Short] = BoxedShortEncoder - - /** @since 3.4.0 */ - implicit val newBoxedBooleanEncoder: Encoder[java.lang.Boolean] = BoxedBooleanEncoder - - // Seqs - private def newSeqEncoder[E](elementEncoder: AgnosticEncoder[E]): AgnosticEncoder[Seq[E]] = { - IterableEncoder( - classTag[Seq[E]], - elementEncoder, - elementEncoder.nullable, - elementEncoder.lenientSerialization) - } - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newIntSeqEncoder: Encoder[Seq[Int]] = newSeqEncoder(PrimitiveIntEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newLongSeqEncoder: Encoder[Seq[Long]] = newSeqEncoder(PrimitiveLongEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newDoubleSeqEncoder: Encoder[Seq[Double]] = newSeqEncoder(PrimitiveDoubleEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newFloatSeqEncoder: Encoder[Seq[Float]] = newSeqEncoder(PrimitiveFloatEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newByteSeqEncoder: Encoder[Seq[Byte]] = newSeqEncoder(PrimitiveByteEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newShortSeqEncoder: Encoder[Seq[Short]] = newSeqEncoder(PrimitiveShortEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newBooleanSeqEncoder: Encoder[Seq[Boolean]] = newSeqEncoder(PrimitiveBooleanEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - val newStringSeqEncoder: Encoder[Seq[String]] = newSeqEncoder(StringEncoder) - - /** - * @since 3.4.0 - * @deprecated - * use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newProductSeqEncoder[A <: Product: TypeTag]: Encoder[Seq[A]] = - newSeqEncoder(ScalaReflection.encoderFor[A]) - - /** @since 3.4.0 */ - implicit def newSequenceEncoder[T <: Seq[_]: TypeTag]: Encoder[T] = - ScalaReflection.encoderFor[T] - - // Maps - /** @since 3.4.0 */ - implicit def newMapEncoder[T <: Map[_, _]: TypeTag]: Encoder[T] = ScalaReflection.encoderFor[T] - - /** - * Notice that we serialize `Set` to Catalyst array. The set property is only kept when - * manipulating the domain objects. The serialization format doesn't keep the set property. When - * we have a Catalyst array which contains duplicated elements and convert it to - * `Dataset[Set[T]]` by using the encoder, the elements will be de-duplicated. - * - * @since 3.4.0 - */ - implicit def newSetEncoder[T <: Set[_]: TypeTag]: Encoder[T] = ScalaReflection.encoderFor[T] - - // Arrays - private def newArrayEncoder[E]( - elementEncoder: AgnosticEncoder[E]): AgnosticEncoder[Array[E]] = { - ArrayEncoder(elementEncoder, elementEncoder.nullable) - } - - /** @since 3.4.0 */ - implicit val newIntArrayEncoder: Encoder[Array[Int]] = newArrayEncoder(PrimitiveIntEncoder) - - /** @since 3.4.0 */ - implicit val newLongArrayEncoder: Encoder[Array[Long]] = newArrayEncoder(PrimitiveLongEncoder) - - /** @since 3.4.0 */ - implicit val newDoubleArrayEncoder: Encoder[Array[Double]] = - newArrayEncoder(PrimitiveDoubleEncoder) - - /** @since 3.4.0 */ - implicit val newFloatArrayEncoder: Encoder[Array[Float]] = newArrayEncoder( - PrimitiveFloatEncoder) - - /** @since 3.4.0 */ - implicit val newByteArrayEncoder: Encoder[Array[Byte]] = BinaryEncoder - - /** @since 3.4.0 */ - implicit val newShortArrayEncoder: Encoder[Array[Short]] = newArrayEncoder( - PrimitiveShortEncoder) - - /** @since 3.4.0 */ - implicit val newBooleanArrayEncoder: Encoder[Array[Boolean]] = - newArrayEncoder(PrimitiveBooleanEncoder) - - /** @since 3.4.0 */ - implicit val newStringArrayEncoder: Encoder[Array[String]] = newArrayEncoder(StringEncoder) - - /** @since 3.4.0 */ - implicit def newProductArrayEncoder[A <: Product: TypeTag]: Encoder[Array[A]] = { - newArrayEncoder(ScalaReflection.encoderFor[A]) - } - - /** - * Creates a [[Dataset]] from a local Seq. - * @since 3.4.0 - */ - implicit def localSeqToDatasetHolder[T: Encoder](s: Seq[T]): DatasetHolder[T] = { - DatasetHolder(session.createDataset(s)) - } -} - -/** - * Lower priority implicit methods for converting Scala objects into [[Dataset]]s. Conflicting - * implicits are placed here to disambiguate resolution. - * - * Reasons for including specific implicits: newProductEncoder - to disambiguate for `List`s which - * are both `Seq` and `Product` - */ -trait LowPrioritySQLImplicits { - - /** @since 3.4.0 */ - implicit def newProductEncoder[T <: Product: TypeTag]: Encoder[T] = - ScalaReflection.encoderFor[T] +/** @inheritdoc */ +abstract class SQLImplicits private[sql] (override val session: SparkSession) + extends api.SQLImplicits { + type DS[U] = Dataset[U] } diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 04f8eeb5c6d46..0663f0186888e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -252,19 +252,8 @@ class SparkSession private[sql] ( lazy val udf: UDFRegistration = new UDFRegistration(this) // scalastyle:off - // Disable style checker so "implicits" object can start with lowercase i - /** - * (Scala-specific) Implicit methods available in Scala for converting common names and Symbols - * into [[Column]]s, and for converting common Scala objects into DataFrame`s. - * - * {{{ - * val sparkSession = SparkSession.builder.getOrCreate() - * import sparkSession.implicits._ - * }}} - * - * @since 3.4.0 - */ - object implicits extends SQLImplicits(this) with Serializable + /** @inheritdoc */ + object implicits extends SQLImplicits(this) // scalastyle:on /** @inheritdoc */ diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index ece4504395f12..972438d0757a7 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -183,6 +183,18 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.streaming.StreamingQueryListener$QueryStartedEvent"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.streaming.StreamingQueryListener$QueryTerminatedEvent"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.streaming.StreamingQueryStatus"), + + // SPARK-49415: Shared SQLImplicits. + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.DatasetHolder"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.DatasetHolder$"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.LowPrioritySQLImplicits"), + ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SQLContext$implicits$"), + ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SQLImplicits"), + ProblemFilters.exclude[IncompatibleResultTypeProblem]("org.apache.spark.sql.SQLImplicits.StringToColumn"), + ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.SQLImplicits.this"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.SQLImplicits$StringToColumn"), + ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SparkSession$implicits$"), + ProblemFilters.exclude[InheritedNewAbstractMethodProblem]("org.apache.spark.sql.SQLImplicits.session"), ) // Default exclude rules diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala b/sql/api/src/main/scala/org/apache/spark/sql/DatasetHolder.scala similarity index 79% rename from sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala rename to sql/api/src/main/scala/org/apache/spark/sql/DatasetHolder.scala index 1c4ffefb897ea..dd7e8e81a088c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/DatasetHolder.scala @@ -18,9 +18,10 @@ package org.apache.spark.sql import org.apache.spark.annotation.Stable +import org.apache.spark.sql.api.Dataset /** - * A container for a [[Dataset]], used for implicit conversions in Scala. + * A container for a [[org.apache.spark.sql.api.Dataset]], used for implicit conversions in Scala. * * To use this, import implicit conversions in SQL: * {{{ @@ -31,15 +32,15 @@ import org.apache.spark.annotation.Stable * @since 1.6.0 */ @Stable -case class DatasetHolder[T] private[sql](private val ds: Dataset[T]) { +class DatasetHolder[T, DS[U] <: Dataset[U]](ds: DS[T]) { // This is declared with parentheses to prevent the Scala compiler from treating // `rdd.toDS("1")` as invoking this toDS and then apply on the returned Dataset. - def toDS(): Dataset[T] = ds + def toDS(): DS[T] = ds // This is declared with parentheses to prevent the Scala compiler from treating // `rdd.toDF("1")` as invoking this toDF and then apply on the returned DataFrame. - def toDF(): DataFrame = ds.toDF() + def toDF(): DS[Row] = ds.toDF().asInstanceOf[DS[Row]] - def toDF(colNames: String*): DataFrame = ds.toDF(colNames : _*) + def toDF(colNames: String*): DS[Row] = ds.toDF(colNames: _*).asInstanceOf[DS[Row]] } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SQLImplicits.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SQLImplicits.scala new file mode 100644 index 0000000000000..f6b44e168390a --- /dev/null +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SQLImplicits.scala @@ -0,0 +1,300 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.api + +import scala.collection.Map +import scala.language.implicitConversions +import scala.reflect.classTag +import scala.reflect.runtime.universe.TypeTag + +import _root_.java + +import org.apache.spark.sql.{ColumnName, DatasetHolder, Encoder, Encoders} +import org.apache.spark.sql.catalyst.ScalaReflection +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{ArrayEncoder, DEFAULT_SCALA_DECIMAL_ENCODER, IterableEncoder, PrimitiveBooleanEncoder, PrimitiveByteEncoder, PrimitiveDoubleEncoder, PrimitiveFloatEncoder, PrimitiveIntEncoder, PrimitiveLongEncoder, PrimitiveShortEncoder, StringEncoder} + +/** + * A collection of implicit methods for converting common Scala objects into + * [[org.apache.spark.sql.api.Dataset]]s. + * + * @since 1.6.0 + */ +abstract class SQLImplicits extends LowPrioritySQLImplicits with Serializable { + type DS[U] <: Dataset[U] + + protected def session: SparkSession + + /** + * Converts $"col name" into a [[org.apache.spark.sql.Column]]. + * + * @since 2.0.0 + */ + implicit class StringToColumn(val sc: StringContext) { + def $(args: Any*): ColumnName = { + new ColumnName(sc.s(args: _*)) + } + } + + // Primitives + + /** @since 1.6.0 */ + implicit def newIntEncoder: Encoder[Int] = Encoders.scalaInt + + /** @since 1.6.0 */ + implicit def newLongEncoder: Encoder[Long] = Encoders.scalaLong + + /** @since 1.6.0 */ + implicit def newDoubleEncoder: Encoder[Double] = Encoders.scalaDouble + + /** @since 1.6.0 */ + implicit def newFloatEncoder: Encoder[Float] = Encoders.scalaFloat + + /** @since 1.6.0 */ + implicit def newByteEncoder: Encoder[Byte] = Encoders.scalaByte + + /** @since 1.6.0 */ + implicit def newShortEncoder: Encoder[Short] = Encoders.scalaShort + + /** @since 1.6.0 */ + implicit def newBooleanEncoder: Encoder[Boolean] = Encoders.scalaBoolean + + /** @since 1.6.0 */ + implicit def newStringEncoder: Encoder[String] = Encoders.STRING + + /** @since 2.2.0 */ + implicit def newJavaDecimalEncoder: Encoder[java.math.BigDecimal] = Encoders.DECIMAL + + /** @since 2.2.0 */ + implicit def newScalaDecimalEncoder: Encoder[scala.math.BigDecimal] = + DEFAULT_SCALA_DECIMAL_ENCODER + + /** @since 2.2.0 */ + implicit def newDateEncoder: Encoder[java.sql.Date] = Encoders.DATE + + /** @since 3.0.0 */ + implicit def newLocalDateEncoder: Encoder[java.time.LocalDate] = Encoders.LOCALDATE + + /** @since 3.4.0 */ + implicit def newLocalDateTimeEncoder: Encoder[java.time.LocalDateTime] = Encoders.LOCALDATETIME + + /** @since 2.2.0 */ + implicit def newTimeStampEncoder: Encoder[java.sql.Timestamp] = Encoders.TIMESTAMP + + /** @since 3.0.0 */ + implicit def newInstantEncoder: Encoder[java.time.Instant] = Encoders.INSTANT + + /** @since 3.2.0 */ + implicit def newDurationEncoder: Encoder[java.time.Duration] = Encoders.DURATION + + /** @since 3.2.0 */ + implicit def newPeriodEncoder: Encoder[java.time.Period] = Encoders.PERIOD + + /** @since 3.2.0 */ + implicit def newJavaEnumEncoder[A <: java.lang.Enum[_]: TypeTag]: Encoder[A] = + ScalaReflection.encoderFor[A] + + // Boxed primitives + + /** @since 2.0.0 */ + implicit def newBoxedIntEncoder: Encoder[java.lang.Integer] = Encoders.INT + + /** @since 2.0.0 */ + implicit def newBoxedLongEncoder: Encoder[java.lang.Long] = Encoders.LONG + + /** @since 2.0.0 */ + implicit def newBoxedDoubleEncoder: Encoder[java.lang.Double] = Encoders.DOUBLE + + /** @since 2.0.0 */ + implicit def newBoxedFloatEncoder: Encoder[java.lang.Float] = Encoders.FLOAT + + /** @since 2.0.0 */ + implicit def newBoxedByteEncoder: Encoder[java.lang.Byte] = Encoders.BYTE + + /** @since 2.0.0 */ + implicit def newBoxedShortEncoder: Encoder[java.lang.Short] = Encoders.SHORT + + /** @since 2.0.0 */ + implicit def newBoxedBooleanEncoder: Encoder[java.lang.Boolean] = Encoders.BOOLEAN + + // Seqs + private def newSeqEncoder[E](elementEncoder: AgnosticEncoder[E]): AgnosticEncoder[Seq[E]] = { + IterableEncoder( + classTag[Seq[E]], + elementEncoder, + elementEncoder.nullable, + elementEncoder.lenientSerialization) + } + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newIntSeqEncoder: Encoder[Seq[Int]] = newSeqEncoder(PrimitiveIntEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newLongSeqEncoder: Encoder[Seq[Long]] = newSeqEncoder(PrimitiveLongEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newDoubleSeqEncoder: Encoder[Seq[Double]] = newSeqEncoder(PrimitiveDoubleEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newFloatSeqEncoder: Encoder[Seq[Float]] = newSeqEncoder(PrimitiveFloatEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newByteSeqEncoder: Encoder[Seq[Byte]] = newSeqEncoder(PrimitiveByteEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newShortSeqEncoder: Encoder[Seq[Short]] = newSeqEncoder(PrimitiveShortEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newBooleanSeqEncoder: Encoder[Seq[Boolean]] = newSeqEncoder(PrimitiveBooleanEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + val newStringSeqEncoder: Encoder[Seq[String]] = newSeqEncoder(StringEncoder) + + /** + * @since 1.6.1 + * @deprecated + * use [[newSequenceEncoder]] + */ + @deprecated("Use newSequenceEncoder instead", "2.2.0") + def newProductSeqEncoder[A <: Product: TypeTag]: Encoder[Seq[A]] = + newSeqEncoder(ScalaReflection.encoderFor[A]) + + /** @since 2.2.0 */ + implicit def newSequenceEncoder[T <: Seq[_]: TypeTag]: Encoder[T] = + ScalaReflection.encoderFor[T] + + // Maps + /** @since 2.3.0 */ + implicit def newMapEncoder[T <: Map[_, _]: TypeTag]: Encoder[T] = ScalaReflection.encoderFor[T] + + /** + * Notice that we serialize `Set` to Catalyst array. The set property is only kept when + * manipulating the domain objects. The serialization format doesn't keep the set property. When + * we have a Catalyst array which contains duplicated elements and convert it to + * `Dataset[Set[T]]` by using the encoder, the elements will be de-duplicated. + * + * @since 2.3.0 + */ + implicit def newSetEncoder[T <: Set[_]: TypeTag]: Encoder[T] = ScalaReflection.encoderFor[T] + + // Arrays + private def newArrayEncoder[E]( + elementEncoder: AgnosticEncoder[E]): AgnosticEncoder[Array[E]] = { + ArrayEncoder(elementEncoder, elementEncoder.nullable) + } + + /** @since 1.6.1 */ + implicit val newIntArrayEncoder: Encoder[Array[Int]] = newArrayEncoder(PrimitiveIntEncoder) + + /** @since 1.6.1 */ + implicit val newLongArrayEncoder: Encoder[Array[Long]] = newArrayEncoder(PrimitiveLongEncoder) + + /** @since 1.6.1 */ + implicit val newDoubleArrayEncoder: Encoder[Array[Double]] = + newArrayEncoder(PrimitiveDoubleEncoder) + + /** @since 1.6.1 */ + implicit val newFloatArrayEncoder: Encoder[Array[Float]] = + newArrayEncoder(PrimitiveFloatEncoder) + + /** @since 1.6.1 */ + implicit val newByteArrayEncoder: Encoder[Array[Byte]] = Encoders.BINARY + + /** @since 1.6.1 */ + implicit val newShortArrayEncoder: Encoder[Array[Short]] = + newArrayEncoder(PrimitiveShortEncoder) + + /** @since 1.6.1 */ + implicit val newBooleanArrayEncoder: Encoder[Array[Boolean]] = + newArrayEncoder(PrimitiveBooleanEncoder) + + /** @since 1.6.1 */ + implicit val newStringArrayEncoder: Encoder[Array[String]] = + newArrayEncoder(StringEncoder) + + /** @since 1.6.1 */ + implicit def newProductArrayEncoder[A <: Product: TypeTag]: Encoder[Array[A]] = + newArrayEncoder(ScalaReflection.encoderFor[A]) + + /** + * Creates a [[Dataset]] from a local Seq. + * @since 1.6.0 + */ + implicit def localSeqToDatasetHolder[T: Encoder](s: Seq[T]): DatasetHolder[T, DS] = { + new DatasetHolder(session.createDataset(s).asInstanceOf[DS[T]]) + } + + /** + * An implicit conversion that turns a Scala `Symbol` into a [[org.apache.spark.sql.Column]]. + * @since 1.3.0 + */ + implicit def symbolToColumn(s: Symbol): ColumnName = new ColumnName(s.name) +} + +/** + * Lower priority implicit methods for converting Scala objects into + * [[org.apache.spark.sql.api.Dataset]]s. Conflicting implicits are placed here to disambiguate + * resolution. + * + * Reasons for including specific implicits: newProductEncoder - to disambiguate for `List`s which + * are both `Seq` and `Product` + */ +trait LowPrioritySQLImplicits { + + /** @since 1.6.0 */ + implicit def newProductEncoder[T <: Product: TypeTag]: Encoder[T] = Encoders.product[T] +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 41d16b16ab1c5..2623db4060ee6 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -505,6 +505,19 @@ abstract class SparkSession extends Serializable with Closeable { */ def read: DataFrameReader + /** + * (Scala-specific) Implicit methods available in Scala for converting common Scala objects into + * `DataFrame`s. + * + * {{{ + * val sparkSession = SparkSession.builder.getOrCreate() + * import sparkSession.implicits._ + * }}} + * + * @since 2.0.0 + */ + val implicits: SQLImplicits + /** * Executes some code block and prints to stdout the time taken to execute the block. This is * available in Scala only and is used primarily for interactive testing and debugging. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala index ffcc0b923f2cb..636899a7acb06 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala @@ -251,8 +251,8 @@ class SQLContext private[sql](val sparkSession: SparkSession) * @group basic * @since 1.3.0 */ - object implicits extends SQLImplicits with Serializable { - protected override def session: SparkSession = self.sparkSession + object implicits extends SQLImplicits { + override protected def session: SparkSession = sparkSession } // scalastyle:on diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala index a657836aafbea..1bc7e3ee98e76 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala @@ -17,259 +17,21 @@ package org.apache.spark.sql -import scala.collection.Map import scala.language.implicitConversions -import scala.reflect.runtime.universe.TypeTag import org.apache.spark.rdd.RDD -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder -/** - * A collection of implicit methods for converting common Scala objects into [[Dataset]]s. - * - * @since 1.6.0 - */ -abstract class SQLImplicits extends LowPrioritySQLImplicits { +/** @inheritdoc */ +abstract class SQLImplicits extends api.SQLImplicits { + type DS[U] = Dataset[U] protected def session: SparkSession - /** - * Converts $"col name" into a [[Column]]. - * - * @since 2.0.0 - */ - implicit class StringToColumn(val sc: StringContext) { - def $(args: Any*): ColumnName = { - new ColumnName(sc.s(args: _*)) - } - } - - // Primitives - - /** @since 1.6.0 */ - implicit def newIntEncoder: Encoder[Int] = Encoders.scalaInt - - /** @since 1.6.0 */ - implicit def newLongEncoder: Encoder[Long] = Encoders.scalaLong - - /** @since 1.6.0 */ - implicit def newDoubleEncoder: Encoder[Double] = Encoders.scalaDouble - - /** @since 1.6.0 */ - implicit def newFloatEncoder: Encoder[Float] = Encoders.scalaFloat - - /** @since 1.6.0 */ - implicit def newByteEncoder: Encoder[Byte] = Encoders.scalaByte - - /** @since 1.6.0 */ - implicit def newShortEncoder: Encoder[Short] = Encoders.scalaShort - - /** @since 1.6.0 */ - implicit def newBooleanEncoder: Encoder[Boolean] = Encoders.scalaBoolean - - /** @since 1.6.0 */ - implicit def newStringEncoder: Encoder[String] = Encoders.STRING - - /** @since 2.2.0 */ - implicit def newJavaDecimalEncoder: Encoder[java.math.BigDecimal] = Encoders.DECIMAL - - /** @since 2.2.0 */ - implicit def newScalaDecimalEncoder: Encoder[scala.math.BigDecimal] = ExpressionEncoder() - - /** @since 2.2.0 */ - implicit def newDateEncoder: Encoder[java.sql.Date] = Encoders.DATE - - /** @since 3.0.0 */ - implicit def newLocalDateEncoder: Encoder[java.time.LocalDate] = Encoders.LOCALDATE - - /** @since 3.4.0 */ - implicit def newLocalDateTimeEncoder: Encoder[java.time.LocalDateTime] = Encoders.LOCALDATETIME - - /** @since 2.2.0 */ - implicit def newTimeStampEncoder: Encoder[java.sql.Timestamp] = Encoders.TIMESTAMP - - /** @since 3.0.0 */ - implicit def newInstantEncoder: Encoder[java.time.Instant] = Encoders.INSTANT - - /** @since 3.2.0 */ - implicit def newDurationEncoder: Encoder[java.time.Duration] = Encoders.DURATION - - /** @since 3.2.0 */ - implicit def newPeriodEncoder: Encoder[java.time.Period] = Encoders.PERIOD - - /** @since 3.2.0 */ - implicit def newJavaEnumEncoder[A <: java.lang.Enum[_] : TypeTag]: Encoder[A] = - ExpressionEncoder() - - // Boxed primitives - - /** @since 2.0.0 */ - implicit def newBoxedIntEncoder: Encoder[java.lang.Integer] = Encoders.INT - - /** @since 2.0.0 */ - implicit def newBoxedLongEncoder: Encoder[java.lang.Long] = Encoders.LONG - - /** @since 2.0.0 */ - implicit def newBoxedDoubleEncoder: Encoder[java.lang.Double] = Encoders.DOUBLE - - /** @since 2.0.0 */ - implicit def newBoxedFloatEncoder: Encoder[java.lang.Float] = Encoders.FLOAT - - /** @since 2.0.0 */ - implicit def newBoxedByteEncoder: Encoder[java.lang.Byte] = Encoders.BYTE - - /** @since 2.0.0 */ - implicit def newBoxedShortEncoder: Encoder[java.lang.Short] = Encoders.SHORT - - /** @since 2.0.0 */ - implicit def newBoxedBooleanEncoder: Encoder[java.lang.Boolean] = Encoders.BOOLEAN - - // Seqs - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newIntSeqEncoder: Encoder[Seq[Int]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newLongSeqEncoder: Encoder[Seq[Long]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newDoubleSeqEncoder: Encoder[Seq[Double]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newFloatSeqEncoder: Encoder[Seq[Float]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newByteSeqEncoder: Encoder[Seq[Byte]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newShortSeqEncoder: Encoder[Seq[Short]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newBooleanSeqEncoder: Encoder[Seq[Boolean]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newStringSeqEncoder: Encoder[Seq[String]] = ExpressionEncoder() - - /** - * @since 1.6.1 - * @deprecated use [[newSequenceEncoder]] - */ - @deprecated("Use newSequenceEncoder instead", "2.2.0") - def newProductSeqEncoder[A <: Product : TypeTag]: Encoder[Seq[A]] = ExpressionEncoder() - - /** @since 2.2.0 */ - implicit def newSequenceEncoder[T <: Seq[_] : TypeTag]: Encoder[T] = ExpressionEncoder() - - // Maps - /** @since 2.3.0 */ - implicit def newMapEncoder[T <: Map[_, _] : TypeTag]: Encoder[T] = ExpressionEncoder() - - /** - * Notice that we serialize `Set` to Catalyst array. The set property is only kept when - * manipulating the domain objects. The serialization format doesn't keep the set property. - * When we have a Catalyst array which contains duplicated elements and convert it to - * `Dataset[Set[T]]` by using the encoder, the elements will be de-duplicated. - * - * @since 2.3.0 - */ - implicit def newSetEncoder[T <: Set[_] : TypeTag]: Encoder[T] = ExpressionEncoder() - - // Arrays - - /** @since 1.6.1 */ - implicit def newIntArrayEncoder: Encoder[Array[Int]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newLongArrayEncoder: Encoder[Array[Long]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newDoubleArrayEncoder: Encoder[Array[Double]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newFloatArrayEncoder: Encoder[Array[Float]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newByteArrayEncoder: Encoder[Array[Byte]] = Encoders.BINARY - - /** @since 1.6.1 */ - implicit def newShortArrayEncoder: Encoder[Array[Short]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newBooleanArrayEncoder: Encoder[Array[Boolean]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newStringArrayEncoder: Encoder[Array[String]] = ExpressionEncoder() - - /** @since 1.6.1 */ - implicit def newProductArrayEncoder[A <: Product : TypeTag]: Encoder[Array[A]] = - ExpressionEncoder() - /** * Creates a [[Dataset]] from an RDD. * * @since 1.6.0 */ - implicit def rddToDatasetHolder[T : Encoder](rdd: RDD[T]): DatasetHolder[T] = { - DatasetHolder(session.createDataset(rdd)) - } - - /** - * Creates a [[Dataset]] from a local Seq. - * @since 1.6.0 - */ - implicit def localSeqToDatasetHolder[T : Encoder](s: Seq[T]): DatasetHolder[T] = { - DatasetHolder(session.createDataset(s)) - } - - /** - * An implicit conversion that turns a Scala `Symbol` into a [[Column]]. - * @since 1.3.0 - */ - implicit def symbolToColumn(s: Symbol): ColumnName = new ColumnName(s.name) - -} - -/** - * Lower priority implicit methods for converting Scala objects into [[Dataset]]s. - * Conflicting implicits are placed here to disambiguate resolution. - * - * Reasons for including specific implicits: - * newProductEncoder - to disambiguate for `List`s which are both `Seq` and `Product` - */ -trait LowPrioritySQLImplicits { - /** @since 1.6.0 */ - implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = Encoders.product[T] - + implicit def rddToDatasetHolder[T : Encoder](rdd: RDD[T]): DatasetHolder[T, Dataset] = + new DatasetHolder(session.createDataset(rdd)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index 137dbaed9f00a..938df206b9792 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -752,19 +752,8 @@ class SparkSession private( // scalastyle:off // Disable style checker so "implicits" object can start with lowercase i - /** - * (Scala-specific) Implicit methods available in Scala for converting - * common Scala objects into `DataFrame`s. - * - * {{{ - * val sparkSession = SparkSession.builder.getOrCreate() - * import sparkSession.implicits._ - * }}} - * - * @since 2.0.0 - */ - object implicits extends SQLImplicits with Serializable { - protected override def session: SparkSession = SparkSession.this + object implicits extends SQLImplicits { + override protected def session: SparkSession = self } // scalastyle:on diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala index 6277f8b459248..8d17edd42442e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala @@ -39,11 +39,6 @@ object typed { // For example, avg in the Scala version returns Scala primitive Double, whose bytecode // signature is just a java.lang.Object; avg in the Java version returns java.lang.Double. - // TODO: This is pretty hacky. Maybe we should have an object for implicit encoders. - private val implicits = new SQLImplicits { - override protected def session: SparkSession = null - } - /** * Average aggregate function. * diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala index d7c00b68828c4..90432dea3a017 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala @@ -35,7 +35,7 @@ private[sql] trait SQLTestData { self => // Helper object to import SQL implicits without a concrete SparkSession private object internalImplicits extends SQLImplicits { - protected override def session: SparkSession = self.spark + override protected def session: SparkSession = self.spark } import internalImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala index 54d6840eb5775..fe5a0f8ee257a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala @@ -240,7 +240,7 @@ private[sql] trait SQLTestUtilsBase * but the implicits import is needed in the constructor. */ protected object testImplicits extends SQLImplicits { - protected override def session: SparkSession = self.spark + override protected def session: SparkSession = self.spark implicit def toRichColumn(c: Column): SparkSession#RichColumn = session.RichColumn(c) } From 94d288e08f2b9b98c2e74a8dcced86b163c1637a Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 24 Sep 2024 08:51:07 +0900 Subject: [PATCH 111/189] [MINOR][PYTHON][DOCS] Fix the docstring of `to_timestamp` ### What changes were proposed in this pull request? Fix the docstring of `to_timestamp` ### Why are the changes needed? `try_to_timestamp` is used in the examples of `to_timestamp` ### Does this PR introduce _any_ user-facing change? doc changes ### How was this patch tested? updated doctests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48207 from zhengruifeng/py_doc_nit_tots. Authored-by: Ruifeng Zheng Signed-off-by: Hyukjin Kwon --- python/pyspark/sql/functions/builtin.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 2d5dbb5946050..2688f9daa23a4 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -9091,15 +9091,19 @@ def to_timestamp(col: "ColumnOrName", format: Optional[str] = None) -> Column: :class:`~pyspark.sql.Column` timestamp value as :class:`pyspark.sql.types.TimestampType` type. + See Also + -------- + :meth:`pyspark.sql.functions.try_to_timestamp` + Examples -------- Example 1: Convert string to a timestamp >>> import pyspark.sql.functions as sf >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) - >>> df.select(sf.try_to_timestamp(df.t).alias('dt')).show() + >>> df.select(sf.to_timestamp(df.t)).show() +-------------------+ - | dt| + | to_timestamp(t)| +-------------------+ |1997-02-28 10:30:00| +-------------------+ @@ -9108,12 +9112,12 @@ def to_timestamp(col: "ColumnOrName", format: Optional[str] = None) -> Column: >>> import pyspark.sql.functions as sf >>> df = spark.createDataFrame([('1997-02-28 10:30:00',)], ['t']) - >>> df.select(sf.try_to_timestamp(df.t, sf.lit('yyyy-MM-dd HH:mm:ss')).alias('dt')).show() - +-------------------+ - | dt| - +-------------------+ - |1997-02-28 10:30:00| - +-------------------+ + >>> df.select(sf.to_timestamp(df.t, 'yyyy-MM-dd HH:mm:ss')).show() + +------------------------------------+ + |to_timestamp(t, yyyy-MM-dd HH:mm:ss)| + +------------------------------------+ + | 1997-02-28 10:30:00| + +------------------------------------+ """ from pyspark.sql.classic.column import _to_java_column @@ -9139,6 +9143,10 @@ def try_to_timestamp(col: "ColumnOrName", format: Optional["ColumnOrName"] = Non format: str, optional format to use to convert timestamp values. + See Also + -------- + :meth:`pyspark.sql.functions.to_timestamp` + Examples -------- Example 1: Convert string to a timestamp From 742265ebb742f9520ca06717be57c6aa2e594191 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Mon, 23 Sep 2024 22:51:54 -0400 Subject: [PATCH 112/189] [SPARK-49429][CONNECT][SQL] Add Shared DataStreamWriter interface ### What changes were proposed in this pull request? This PR adds a shared DataStreamWriter to sql. ### Why are the changes needed? We are creating a unified Scala interface for sql. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48212 from hvanhovell/SPARK-49429. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../scala/org/apache/spark/sql/Dataset.scala | 7 +- .../sql/streaming/DataStreamWriter.scala | 252 +++---------- .../spark/sql/api/DataStreamWriter.scala | 193 ++++++++++ .../org/apache/spark/sql/api/Dataset.scala | 8 + .../scala/org/apache/spark/sql/Dataset.scala | 7 +- .../sql/streaming/DataStreamWriter.scala | 343 +++++------------- .../sql/streaming/StreamingQueryManager.scala | 6 +- 7 files changed, 340 insertions(+), 476 deletions(-) create mode 100644 sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamWriter.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index accfff9f2b073..d2877ccaf06c9 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -1035,12 +1035,7 @@ class Dataset[T] private[sql] ( new MergeIntoWriterImpl[T](table, this, condition) } - /** - * Interface for saving the content of the streaming Dataset out into external storage. - * - * @group basic - * @since 3.5.0 - */ + /** @inheritdoc */ def writeStream: DataStreamWriter[T] = { new DataStreamWriter[T](this) } diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala index c8c714047788b..9fcc31e562682 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala @@ -29,9 +29,8 @@ import org.apache.spark.api.java.function.VoidFunction2 import org.apache.spark.connect.proto import org.apache.spark.connect.proto.Command import org.apache.spark.connect.proto.WriteStreamOperationStart -import org.apache.spark.internal.Logging -import org.apache.spark.sql.{Dataset, ForeachWriter} -import org.apache.spark.sql.connect.common.{DataTypeProtoConverter, ForeachWriterPacket, UdfUtils} +import org.apache.spark.sql.{api, Dataset, ForeachWriter} +import org.apache.spark.sql.connect.common.{DataTypeProtoConverter, ForeachWriterPacket} import org.apache.spark.sql.execution.streaming.AvailableNowTrigger import org.apache.spark.sql.execution.streaming.ContinuousTrigger import org.apache.spark.sql.execution.streaming.OneTimeTrigger @@ -47,63 +46,23 @@ import org.apache.spark.util.SparkSerDeUtils * @since 3.5.0 */ @Evolving -final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { +final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends api.DataStreamWriter[T] { + override type DS[U] = Dataset[U] - /** - * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink.

    • - * `OutputMode.Append()`: only the new rows in the streaming DataFrame/Dataset will be written - * to the sink.
    • `OutputMode.Complete()`: all the rows in the streaming - * DataFrame/Dataset will be written to the sink every time there are some updates.
    • - * `OutputMode.Update()`: only the rows that were updated in the streaming DataFrame/Dataset - * will be written to the sink every time there are some updates. If the query doesn't contain - * aggregations, it will be equivalent to `OutputMode.Append()` mode.
    - * - * @since 3.5.0 - */ - def outputMode(outputMode: OutputMode): DataStreamWriter[T] = { + /** @inheritdoc */ + def outputMode(outputMode: OutputMode): this.type = { sinkBuilder.setOutputMode(outputMode.toString.toLowerCase(Locale.ROOT)) this } - /** - * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink.
    • - * `append`: only the new rows in the streaming DataFrame/Dataset will be written to the - * sink.
    • `complete`: all the rows in the streaming DataFrame/Dataset will be written - * to the sink every time there are some updates.
    • `update`: only the rows that were - * updated in the streaming DataFrame/Dataset will be written to the sink every time there are - * some updates. If the query doesn't contain aggregations, it will be equivalent to `append` - * mode.
    - * - * @since 3.5.0 - */ - def outputMode(outputMode: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def outputMode(outputMode: String): this.type = { sinkBuilder.setOutputMode(outputMode) this } - /** - * Set the trigger for the stream query. The default value is `ProcessingTime(0)` and it will - * run the query as fast as possible. - * - * Scala Example: - * {{{ - * df.writeStream.trigger(ProcessingTime("10 seconds")) - * - * import scala.concurrent.duration._ - * df.writeStream.trigger(ProcessingTime(10.seconds)) - * }}} - * - * Java Example: - * {{{ - * df.writeStream().trigger(ProcessingTime.create("10 seconds")) - * - * import java.util.concurrent.TimeUnit - * df.writeStream().trigger(ProcessingTime.create(10, TimeUnit.SECONDS)) - * }}} - * - * @since 3.5.0 - */ - def trigger(trigger: Trigger): DataStreamWriter[T] = { + /** @inheritdoc */ + def trigger(trigger: Trigger): this.type = { trigger match { case ProcessingTimeTrigger(intervalMs) => sinkBuilder.setProcessingTimeInterval(s"$intervalMs milliseconds") @@ -117,123 +76,54 @@ final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { this } - /** - * Specifies the name of the [[StreamingQuery]] that can be started with `start()`. This name - * must be unique among all the currently active queries in the associated SQLContext. - * - * @since 3.5.0 - */ - def queryName(queryName: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def queryName(queryName: String): this.type = { sinkBuilder.setQueryName(queryName) this } - /** - * Specifies the underlying output data source. - * - * @since 3.5.0 - */ - def format(source: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def format(source: String): this.type = { sinkBuilder.setFormat(source) this } - /** - * Partitions the output by the given columns on the file system. If specified, the output is - * laid out on the file system similar to Hive's partitioning scheme. As an example, when we - * partition a dataset by year and then month, the directory layout would look like: - * - *
    • year=2016/month=01/
    • year=2016/month=02/
    - * - * Partitioning is one of the most widely used techniques to optimize physical data layout. It - * provides a coarse-grained index for skipping unnecessary data reads when queries have - * predicates on the partitioned columns. In order for partitioning to work well, the number of - * distinct values in each column should typically be less than tens of thousands. - * - * @since 3.5.0 - */ + /** @inheritdoc */ @scala.annotation.varargs - def partitionBy(colNames: String*): DataStreamWriter[T] = { + def partitionBy(colNames: String*): this.type = { sinkBuilder.clearPartitioningColumnNames() sinkBuilder.addAllPartitioningColumnNames(colNames.asJava) this } - /** - * Clusters the output by the given columns. If specified, the output is laid out such that - * records with similar values on the clustering column are grouped together in the same file. - * - * Clustering improves query efficiency by allowing queries with predicates on the clustering - * columns to skip unnecessary data. Unlike partitioning, clustering can be used on very high - * cardinality columns. - * - * @since 4.0.0 - */ + /** @inheritdoc */ @scala.annotation.varargs - def clusterBy(colNames: String*): DataStreamWriter[T] = { + def clusterBy(colNames: String*): this.type = { sinkBuilder.clearClusteringColumnNames() sinkBuilder.addAllClusteringColumnNames(colNames.asJava) this } - /** - * Adds an output option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def option(key: String, value: String): this.type = { sinkBuilder.putOptions(key, value) this } - /** - * Adds an output option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Boolean): DataStreamWriter[T] = option(key, value.toString) - - /** - * Adds an output option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Long): DataStreamWriter[T] = option(key, value.toString) - - /** - * Adds an output option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Double): DataStreamWriter[T] = option(key, value.toString) - - /** - * (Scala-specific) Adds output options for the underlying data source. - * - * @since 3.5.0 - */ - def options(options: scala.collection.Map[String, String]): DataStreamWriter[T] = { + /** @inheritdoc */ + def options(options: scala.collection.Map[String, String]): this.type = { this.options(options.asJava) this } - /** - * Adds output options for the underlying data source. - * - * @since 3.5.0 - */ - def options(options: java.util.Map[String, String]): DataStreamWriter[T] = { + /** @inheritdoc */ + def options(options: java.util.Map[String, String]): this.type = { sinkBuilder.putAllOptions(options) this } - /** - * Sets the output of the streaming query to be processed using the provided writer object. - * object. See [[org.apache.spark.sql.ForeachWriter]] for more details on the lifecycle and - * semantics. - * @since 3.5.0 - */ - def foreach(writer: ForeachWriter[T]): DataStreamWriter[T] = { + /** @inheritdoc */ + def foreach(writer: ForeachWriter[T]): this.type = { val serialized = SparkSerDeUtils.serialize(ForeachWriterPacket(writer, ds.agnosticEncoder)) val scalaWriterBuilder = proto.ScalarScalaUDF .newBuilder() @@ -242,21 +132,9 @@ final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { this } - /** - * :: Experimental :: - * - * (Scala-specific) Sets the output of the streaming query to be processed using the provided - * function. This is supported only in the micro-batch execution modes (that is, when the - * trigger is not continuous). In every micro-batch, the provided function will be called in - * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. The - * batchId can be used to deduplicate and transactionally write the output (that is, the - * provided Dataset) to external systems. The output Dataset is guaranteed to be exactly the - * same for the same batchId (assuming all operations are deterministic in the query). - * - * @since 3.5.0 - */ + /** @inheritdoc */ @Evolving - def foreachBatch(function: (Dataset[T], Long) => Unit): DataStreamWriter[T] = { + def foreachBatch(function: (Dataset[T], Long) => Unit): this.type = { val serializedFn = SparkSerDeUtils.serialize(function) sinkBuilder.getForeachBatchBuilder.getScalaFunctionBuilder .setPayload(ByteString.copyFrom(serializedFn)) @@ -265,48 +143,13 @@ final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { this } - /** - * :: Experimental :: - * - * (Java-specific) Sets the output of the streaming query to be processed using the provided - * function. This is supported only in the micro-batch execution modes (that is, when the - * trigger is not continuous). In every micro-batch, the provided function will be called in - * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. The - * batchId can be used to deduplicate and transactionally write the output (that is, the - * provided Dataset) to external systems. The output Dataset is guaranteed to be exactly the - * same for the same batchId (assuming all operations are deterministic in the query). - * - * @since 3.5.0 - */ - @Evolving - def foreachBatch(function: VoidFunction2[Dataset[T], java.lang.Long]): DataStreamWriter[T] = { - foreachBatch(UdfUtils.foreachBatchFuncToScalaFunc(function)) - } - - /** - * Starts the execution of the streaming query, which will continually output results to the - * given path as new data arrives. The returned [[StreamingQuery]] object can be used to - * interact with the stream. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def start(path: String): StreamingQuery = { sinkBuilder.setPath(path) start() } - /** - * Starts the execution of the streaming query, which will continually output results to the - * given path as new data arrives. The returned [[StreamingQuery]] object can be used to - * interact with the stream. Throws a `TimeoutException` if the following conditions are met: - * - Another run of the same streaming query, that is a streaming query sharing the same - * checkpoint location, is already active on the same Spark Driver - * - The SQL configuration `spark.sql.streaming.stopActiveRunOnRestart` is enabled - * - The active run cannot be stopped within the timeout controlled by the SQL configuration - * `spark.sql.streaming.stopTimeout` - * - * @since 3.5.0 - */ + /** @inheritdoc */ @throws[TimeoutException] def start(): StreamingQuery = { val startCmd = Command @@ -323,22 +166,7 @@ final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { RemoteStreamingQuery.fromStartCommandResponse(ds.sparkSession, resp) } - /** - * Starts the execution of the streaming query, which will continually output results to the - * given table as new data arrives. The returned [[StreamingQuery]] object can be used to - * interact with the stream. - * - * For v1 table, partitioning columns provided by `partitionBy` will be respected no matter the - * table exists or not. A new table will be created if the table not exists. - * - * For v2 table, `partitionBy` will be ignored if the table already exists. `partitionBy` will - * be respected only if the v2 table does not exist. Besides, the v2 table created by this API - * lacks some functionalities (e.g., customized properties, options, and serde info). If you - * need them, please create the v2 table manually before the execution to avoid creating a table - * with incomplete information. - * - * @since 3.5.0 - */ + /** @inheritdoc */ @Evolving @throws[TimeoutException] def toTable(tableName: String): StreamingQuery = { @@ -346,6 +174,24 @@ final class DataStreamWriter[T] private[sql] (ds: Dataset[T]) extends Logging { start() } + /////////////////////////////////////////////////////////////////////////////////////// + // Covariant Overrides + /////////////////////////////////////////////////////////////////////////////////////// + + /** @inheritdoc */ + override def option(key: String, value: Boolean): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Long): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Double): this.type = super.option(key, value) + + /** @inheritdoc */ + @Evolving + override def foreachBatch(function: VoidFunction2[Dataset[T], java.lang.Long]): this.type = + super.foreachBatch(function) + private val sinkBuilder = WriteStreamOperationStart .newBuilder() .setInput(ds.plan.getRoot) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamWriter.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamWriter.scala new file mode 100644 index 0000000000000..7762708e9520c --- /dev/null +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamWriter.scala @@ -0,0 +1,193 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.api + +import _root_.java +import _root_.java.util.concurrent.TimeoutException + +import org.apache.spark.annotation.Evolving +import org.apache.spark.api.java.function.VoidFunction2 +import org.apache.spark.sql.{ForeachWriter, WriteConfigMethods} +import org.apache.spark.sql.streaming.{OutputMode, Trigger} + +/** + * Interface used to write a streaming `Dataset` to external storage systems (e.g. file systems, + * key-value stores, etc). Use `Dataset.writeStream` to access this. + * + * @since 2.0.0 + */ +@Evolving +abstract class DataStreamWriter[T] extends WriteConfigMethods[DataStreamWriter[T]] { + type DS[U] <: Dataset[U] + + /** + * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink.
    • + * `OutputMode.Append()`: only the new rows in the streaming DataFrame/Dataset will be written + * to the sink.
    • `OutputMode.Complete()`: all the rows in the streaming + * DataFrame/Dataset will be written to the sink every time there are some updates.
    • + * `OutputMode.Update()`: only the rows that were updated in the streaming DataFrame/Dataset + * will be written to the sink every time there are some updates. If the query doesn't contain + * aggregations, it will be equivalent to `OutputMode.Append()` mode.
    + * + * @since 2.0.0 + */ + def outputMode(outputMode: OutputMode): this.type + + /** + * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink.
    • + * `append`: only the new rows in the streaming DataFrame/Dataset will be written to the + * sink.
    • `complete`: all the rows in the streaming DataFrame/Dataset will be written + * to the sink every time there are some updates.
    • `update`: only the rows that were + * updated in the streaming DataFrame/Dataset will be written to the sink every time there are + * some updates. If the query doesn't contain aggregations, it will be equivalent to `append` + * mode.
    + * + * @since 2.0.0 + */ + def outputMode(outputMode: String): this.type + + /** + * Set the trigger for the stream query. The default value is `ProcessingTime(0)` and it will + * run the query as fast as possible. + * + * Scala Example: + * {{{ + * df.writeStream.trigger(ProcessingTime("10 seconds")) + * + * import scala.concurrent.duration._ + * df.writeStream.trigger(ProcessingTime(10.seconds)) + * }}} + * + * Java Example: + * {{{ + * df.writeStream().trigger(ProcessingTime.create("10 seconds")) + * + * import java.util.concurrent.TimeUnit + * df.writeStream().trigger(ProcessingTime.create(10, TimeUnit.SECONDS)) + * }}} + * + * @since 2.0.0 + */ + def trigger(trigger: Trigger): this.type + + /** + * Specifies the name of the [[org.apache.spark.sql.api.StreamingQuery]] that can be started + * with `start()`. This name must be unique among all the currently active queries in the + * associated SparkSession. + * + * @since 2.0.0 + */ + def queryName(queryName: String): this.type + + /** + * Sets the output of the streaming query to be processed using the provided writer object. + * object. See [[org.apache.spark.sql.ForeachWriter]] for more details on the lifecycle and + * semantics. + * + * @since 2.0.0 + */ + def foreach(writer: ForeachWriter[T]): this.type + + /** + * :: Experimental :: + * + * (Scala-specific) Sets the output of the streaming query to be processed using the provided + * function. This is supported only in the micro-batch execution modes (that is, when the + * trigger is not continuous). In every micro-batch, the provided function will be called in + * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. The + * batchId can be used to deduplicate and transactionally write the output (that is, the + * provided Dataset) to external systems. The output Dataset is guaranteed to be exactly the + * same for the same batchId (assuming all operations are deterministic in the query). + * + * @since 2.4.0 + */ + @Evolving + def foreachBatch(function: (DS[T], Long) => Unit): this.type + + /** + * :: Experimental :: + * + * (Java-specific) Sets the output of the streaming query to be processed using the provided + * function. This is supported only in the micro-batch execution modes (that is, when the + * trigger is not continuous). In every micro-batch, the provided function will be called in + * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. The + * batchId can be used to deduplicate and transactionally write the output (that is, the + * provided Dataset) to external systems. The output Dataset is guaranteed to be exactly the + * same for the same batchId (assuming all operations are deterministic in the query). + * + * @since 2.4.0 + */ + @Evolving + def foreachBatch(function: VoidFunction2[DS[T], java.lang.Long]): this.type = { + foreachBatch((batchDs: DS[T], batchId: Long) => function.call(batchDs, batchId)) + } + + /** + * Starts the execution of the streaming query, which will continually output results to the + * given path as new data arrives. The returned [[org.apache.spark.sql.api.StreamingQuery]] + * object can be used to interact with the stream. + * + * @since 2.0.0 + */ + def start(path: String): StreamingQuery + + /** + * Starts the execution of the streaming query, which will continually output results to the + * given path as new data arrives. The returned [[org.apache.spark.sql.api.StreamingQuery]] + * object can be used to interact with the stream. Throws a `TimeoutException` if the following + * conditions are met: + * - Another run of the same streaming query, that is a streaming query sharing the same + * checkpoint location, is already active on the same Spark Driver + * - The SQL configuration `spark.sql.streaming.stopActiveRunOnRestart` is enabled + * - The active run cannot be stopped within the timeout controlled by the SQL configuration + * `spark.sql.streaming.stopTimeout` + * + * @since 2.0.0 + */ + @throws[TimeoutException] + def start(): StreamingQuery + + /** + * Starts the execution of the streaming query, which will continually output results to the + * given table as new data arrives. The returned [[org.apache.spark.sql.api.StreamingQuery]] + * object can be used to interact with the stream. + * + * For v1 table, partitioning columns provided by `partitionBy` will be respected no matter the + * table exists or not. A new table will be created if the table not exists. + * + * For v2 table, `partitionBy` will be ignored if the table already exists. `partitionBy` will + * be respected only if the v2 table does not exist. Besides, the v2 table created by this API + * lacks some functionalities (e.g., customized properties, options, and serde info). If you + * need them, please create the v2 table manually before the execution to avoid creating a table + * with incomplete information. + * + * @since 3.1.0 + */ + @Evolving + @throws[TimeoutException] + def toTable(tableName: String): StreamingQuery + + /////////////////////////////////////////////////////////////////////////////////////// + // Covariant Overrides + /////////////////////////////////////////////////////////////////////////////////////// + override def option(key: String, value: Boolean): this.type = + super.option(key, value).asInstanceOf[this.type] + override def option(key: String, value: Long): this.type = + super.option(key, value).asInstanceOf[this.type] + override def option(key: String, value: Double): this.type = + super.option(key, value).asInstanceOf[this.type] +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala index 6eef034aa5157..06a6148a7c188 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala @@ -3017,6 +3017,14 @@ abstract class Dataset[T] extends Serializable { */ def mergeInto(table: String, condition: Column): MergeIntoWriter[T] + /** + * Interface for saving the content of the streaming Dataset out into external storage. + * + * @group basic + * @since 2.0.0 + */ + def writeStream: DataStreamWriter[T] + /** * Create a write configuration builder for v2 sources. * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index ef628ca612b49..80ec70a7864c3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -1618,12 +1618,7 @@ class Dataset[T] private[sql]( new MergeIntoWriterImpl[T](table, this, condition) } - /** - * Interface for saving the content of the streaming Dataset out into external storage. - * - * @group basic - * @since 2.0.0 - */ + /** @inheritdoc */ def writeStream: DataStreamWriter[T] = { if (!isStreaming) { logicalPlan.failAnalysis( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala index ab4d350c1e68c..b0233d2c51b75 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamWriter.scala @@ -55,253 +55,101 @@ import org.apache.spark.util.Utils * @since 2.0.0 */ @Evolving -final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { - import DataStreamWriter._ +final class DataStreamWriter[T] private[sql](ds: Dataset[T]) extends api.DataStreamWriter[T] { + type DS[U] = Dataset[U] - private val df = ds.toDF() - - /** - * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink. - *
      - *
    • `OutputMode.Append()`: only the new rows in the streaming DataFrame/Dataset will be - * written to the sink.
      • - *
    • `OutputMode.Complete()`: all the rows in the streaming DataFrame/Dataset will be written - * to the sink every time there are some updates.
      • - *
    • `OutputMode.Update()`: only the rows that were updated in the streaming - * DataFrame/Dataset will be written to the sink every time there are some updates. - * If the query doesn't contain aggregations, it will be equivalent to - * `OutputMode.Append()` mode.
      • - *
    - * - * @since 2.0.0 - */ - def outputMode(outputMode: OutputMode): DataStreamWriter[T] = { + /** @inheritdoc */ + def outputMode(outputMode: OutputMode): this.type = { this.outputMode = outputMode this } - /** - * Specifies how data of a streaming DataFrame/Dataset is written to a streaming sink. - *
      - *
    • `append`: only the new rows in the streaming DataFrame/Dataset will be written to - * the sink.
      • - *
    • `complete`: all the rows in the streaming DataFrame/Dataset will be written to the sink - * every time there are some updates.
      • - *
    • `update`: only the rows that were updated in the streaming DataFrame/Dataset will - * be written to the sink every time there are some updates. If the query doesn't - * contain aggregations, it will be equivalent to `append` mode.
      • - *
    - * - * @since 2.0.0 - */ - def outputMode(outputMode: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def outputMode(outputMode: String): this.type = { this.outputMode = InternalOutputModes(outputMode) this } - /** - * Set the trigger for the stream query. The default value is `ProcessingTime(0)` and it will run - * the query as fast as possible. - * - * Scala Example: - * {{{ - * df.writeStream.trigger(ProcessingTime("10 seconds")) - * - * import scala.concurrent.duration._ - * df.writeStream.trigger(ProcessingTime(10.seconds)) - * }}} - * - * Java Example: - * {{{ - * df.writeStream().trigger(ProcessingTime.create("10 seconds")) - * - * import java.util.concurrent.TimeUnit - * df.writeStream().trigger(ProcessingTime.create(10, TimeUnit.SECONDS)) - * }}} - * - * @since 2.0.0 - */ - def trigger(trigger: Trigger): DataStreamWriter[T] = { + /** @inheritdoc */ + def trigger(trigger: Trigger): this.type = { this.trigger = trigger this } - /** - * Specifies the name of the [[StreamingQuery]] that can be started with `start()`. - * This name must be unique among all the currently active queries in the associated SQLContext. - * - * @since 2.0.0 - */ - def queryName(queryName: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def queryName(queryName: String): this.type = { this.extraOptions += ("queryName" -> queryName) this } - /** - * Specifies the underlying output data source. - * - * @since 2.0.0 - */ - def format(source: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def format(source: String): this.type = { this.source = source this } - /** - * Partitions the output by the given columns on the file system. If specified, the output is - * laid out on the file system similar to Hive's partitioning scheme. As an example, when we - * partition a dataset by year and then month, the directory layout would look like: - * - *
      - *
    • year=2016/month=01/
      • - *
    • year=2016/month=02/
      • - *
    - * - * Partitioning is one of the most widely used techniques to optimize physical data layout. - * It provides a coarse-grained index for skipping unnecessary data reads when queries have - * predicates on the partitioned columns. In order for partitioning to work well, the number - * of distinct values in each column should typically be less than tens of thousands. - * - * @since 2.0.0 - */ + /** @inheritdoc */ @scala.annotation.varargs - def partitionBy(colNames: String*): DataStreamWriter[T] = { + def partitionBy(colNames: String*): this.type = { this.partitioningColumns = Option(colNames) validatePartitioningAndClustering() this } - /** - * Clusters the output by the given columns. If specified, the output is laid out such that - * records with similar values on the clustering column are grouped together in the same file. - * - * Clustering improves query efficiency by allowing queries with predicates on the clustering - * columns to skip unnecessary data. Unlike partitioning, clustering can be used on very high - * cardinality columns. - * - * @since 4.0.0 - */ + /** @inheritdoc */ @scala.annotation.varargs - def clusterBy(colNames: String*): DataStreamWriter[T] = { + def clusterBy(colNames: String*): this.type = { this.clusteringColumns = Option(colNames) validatePartitioningAndClustering() this } - /** - * Adds an output option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: String): DataStreamWriter[T] = { + /** @inheritdoc */ + def option(key: String, value: String): this.type = { this.extraOptions += (key -> value) this } - /** - * Adds an output option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Boolean): DataStreamWriter[T] = option(key, value.toString) - - /** - * Adds an output option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Long): DataStreamWriter[T] = option(key, value.toString) - - /** - * Adds an output option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Double): DataStreamWriter[T] = option(key, value.toString) - - /** - * (Scala-specific) Adds output options for the underlying data source. - * - * @since 2.0.0 - */ - def options(options: scala.collection.Map[String, String]): DataStreamWriter[T] = { + /** @inheritdoc */ + def options(options: scala.collection.Map[String, String]): this.type = { this.extraOptions ++= options this } - /** - * Adds output options for the underlying data source. - * - * @since 2.0.0 - */ - def options(options: java.util.Map[String, String]): DataStreamWriter[T] = { + /** @inheritdoc */ + def options(options: java.util.Map[String, String]): this.type = { this.options(options.asScala) this } - /** - * Starts the execution of the streaming query, which will continually output results to the given - * path as new data arrives. The returned [[StreamingQuery]] object can be used to interact with - * the stream. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def start(path: String): StreamingQuery = { - if (!df.sparkSession.sessionState.conf.legacyPathOptionBehavior && + if (!ds.sparkSession.sessionState.conf.legacyPathOptionBehavior && extraOptions.contains("path")) { throw QueryCompilationErrors.setPathOptionAndCallWithPathParameterError("start") } startInternal(Some(path)) } - /** - * Starts the execution of the streaming query, which will continually output results to the given - * path as new data arrives. The returned [[StreamingQuery]] object can be used to interact with - * the stream. Throws a `TimeoutException` if the following conditions are met: - * - Another run of the same streaming query, that is a streaming query - * sharing the same checkpoint location, is already active on the same - * Spark Driver - * - The SQL configuration `spark.sql.streaming.stopActiveRunOnRestart` - * is enabled - * - The active run cannot be stopped within the timeout controlled by - * the SQL configuration `spark.sql.streaming.stopTimeout` - * - * @since 2.0.0 - */ + /** @inheritdoc */ @throws[TimeoutException] def start(): StreamingQuery = startInternal(None) - /** - * Starts the execution of the streaming query, which will continually output results to the given - * table as new data arrives. The returned [[StreamingQuery]] object can be used to interact with - * the stream. - * - * For v1 table, partitioning columns provided by `partitionBy` will be respected no matter the - * table exists or not. A new table will be created if the table not exists. - * - * For v2 table, `partitionBy` will be ignored if the table already exists. `partitionBy` will be - * respected only if the v2 table does not exist. Besides, the v2 table created by this API lacks - * some functionalities (e.g., customized properties, options, and serde info). If you need them, - * please create the v2 table manually before the execution to avoid creating a table with - * incomplete information. - * - * @since 3.1.0 - */ + /** @inheritdoc */ @Evolving @throws[TimeoutException] def toTable(tableName: String): StreamingQuery = { - this.tableName = tableName - import df.sparkSession.sessionState.analyzer.CatalogAndIdentifier + import ds.sparkSession.sessionState.analyzer.CatalogAndIdentifier import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ - val parser = df.sparkSession.sessionState.sqlParser + val parser = ds.sparkSession.sessionState.sqlParser val originalMultipartIdentifier = parser.parseMultipartIdentifier(tableName) val CatalogAndIdentifier(catalog, identifier) = originalMultipartIdentifier // Currently we don't create a logical streaming writer node in logical plan, so cannot rely // on analyzer to resolve it. Directly lookup only for temp view to provide clearer message. // TODO (SPARK-27484): we should add the writing node before the plan is analyzed. - if (df.sparkSession.sessionState.catalog.isTempView(originalMultipartIdentifier)) { + if (ds.sparkSession.sessionState.catalog.isTempView(originalMultipartIdentifier)) { throw QueryCompilationErrors.tempViewNotSupportStreamingWriteError(tableName) } @@ -327,14 +175,14 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { extraOptions.get("path"), None, None, - false) + external = false) val cmd = CreateTable( UnresolvedIdentifier(originalMultipartIdentifier), - df.schema.asNullable.map(ColumnDefinition.fromV1Column(_, parser)), + ds.schema.asNullable.map(ColumnDefinition.fromV1Column(_, parser)), partitioningOrClusteringTransform, tableSpec, ignoreIfExists = false) - Dataset.ofRows(df.sparkSession, cmd) + Dataset.ofRows(ds.sparkSession, cmd) } val tableInstance = catalog.asTableCatalog.loadTable(identifier) @@ -371,34 +219,34 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { throw QueryCompilationErrors.cannotOperateOnHiveDataSourceFilesError("write") } - if (source == SOURCE_NAME_MEMORY) { - assertNotPartitioned(SOURCE_NAME_MEMORY) + if (source == DataStreamWriter.SOURCE_NAME_MEMORY) { + assertNotPartitioned(DataStreamWriter.SOURCE_NAME_MEMORY) if (extraOptions.get("queryName").isEmpty) { throw QueryCompilationErrors.queryNameNotSpecifiedForMemorySinkError() } val sink = new MemorySink() - val resultDf = Dataset.ofRows(df.sparkSession, - MemoryPlan(sink, DataTypeUtils.toAttributes(df.schema))) + val resultDf = Dataset.ofRows(ds.sparkSession, + MemoryPlan(sink, DataTypeUtils.toAttributes(ds.schema))) val recoverFromCheckpoint = outputMode == OutputMode.Complete() val query = startQuery(sink, extraOptions, recoverFromCheckpoint = recoverFromCheckpoint, catalogTable = catalogTable) resultDf.createOrReplaceTempView(query.name) query - } else if (source == SOURCE_NAME_FOREACH) { - assertNotPartitioned(SOURCE_NAME_FOREACH) + } else if (source == DataStreamWriter.SOURCE_NAME_FOREACH) { + assertNotPartitioned(DataStreamWriter.SOURCE_NAME_FOREACH) val sink = ForeachWriterTable[Any](foreachWriter, foreachWriterEncoder) startQuery(sink, extraOptions, catalogTable = catalogTable) - } else if (source == SOURCE_NAME_FOREACH_BATCH) { - assertNotPartitioned(SOURCE_NAME_FOREACH_BATCH) + } else if (source == DataStreamWriter.SOURCE_NAME_FOREACH_BATCH) { + assertNotPartitioned(DataStreamWriter.SOURCE_NAME_FOREACH_BATCH) if (trigger.isInstanceOf[ContinuousTrigger]) { throw QueryCompilationErrors.sourceNotSupportedWithContinuousTriggerError(source) } val sink = new ForeachBatchSink[T](foreachBatchWriter, ds.exprEnc) startQuery(sink, extraOptions, catalogTable = catalogTable) } else { - val cls = DataSource.lookupDataSource(source, df.sparkSession.sessionState.conf) + val cls = DataSource.lookupDataSource(source, ds.sparkSession.sessionState.conf) val disabledSources = - Utils.stringToSeq(df.sparkSession.sessionState.conf.disabledV2StreamingWriters) + Utils.stringToSeq(ds.sparkSession.sessionState.conf.disabledV2StreamingWriters) val useV1Source = disabledSources.contains(cls.getCanonicalName) || // file source v2 does not support streaming yet. classOf[FileDataSourceV2].isAssignableFrom(cls) @@ -412,7 +260,7 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { val sink = if (classOf[TableProvider].isAssignableFrom(cls) && !useV1Source) { val provider = cls.getConstructor().newInstance().asInstanceOf[TableProvider] val sessionOptions = DataSourceV2Utils.extractSessionConfigs( - source = provider, conf = df.sparkSession.sessionState.conf) + source = provider, conf = ds.sparkSession.sessionState.conf) val finalOptions = sessionOptions.filter { case (k, _) => !optionsWithPath.contains(k) } ++ optionsWithPath.originalMap val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) @@ -420,7 +268,7 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { // to `getTable`. This is for avoiding schema inference, which can be very expensive. // If the query schema is not compatible with the existing data, the behavior is undefined. val outputSchema = if (provider.supportsExternalMetadata()) { - Some(df.schema) + Some(ds.schema) } else { None } @@ -450,12 +298,12 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { recoverFromCheckpoint: Boolean = true, catalogAndIdent: Option[(TableCatalog, Identifier)] = None, catalogTable: Option[CatalogTable] = None): StreamingQuery = { - val useTempCheckpointLocation = SOURCES_ALLOW_ONE_TIME_QUERY.contains(source) + val useTempCheckpointLocation = DataStreamWriter.SOURCES_ALLOW_ONE_TIME_QUERY.contains(source) - df.sparkSession.sessionState.streamingQueryManager.startQuery( + ds.sparkSession.sessionState.streamingQueryManager.startQuery( newOptions.get("queryName"), newOptions.get("checkpointLocation"), - df, + ds, newOptions.originalMap, sink, outputMode, @@ -480,26 +328,21 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { case None => optionsWithoutClusteringKey } val ds = DataSource( - df.sparkSession, + this.ds.sparkSession, className = source, options = optionsWithClusteringColumns, partitionColumns = normalizedParCols.getOrElse(Nil)) ds.createSink(outputMode) } - /** - * Sets the output of the streaming query to be processed using the provided writer object. - * object. See [[org.apache.spark.sql.ForeachWriter]] for more details on the lifecycle and - * semantics. - * @since 2.0.0 - */ - def foreach(writer: ForeachWriter[T]): DataStreamWriter[T] = { + /** @inheritdoc */ + def foreach(writer: ForeachWriter[T]): this.type = { foreachImplementation(writer.asInstanceOf[ForeachWriter[Any]]) } private[sql] def foreachImplementation(writer: ForeachWriter[Any], - encoder: Option[ExpressionEncoder[Any]] = None): DataStreamWriter[T] = { - this.source = SOURCE_NAME_FOREACH + encoder: Option[ExpressionEncoder[Any]] = None): this.type = { + this.source = DataStreamWriter.SOURCE_NAME_FOREACH this.foreachWriter = if (writer != null) { ds.sparkSession.sparkContext.clean(writer) } else { @@ -509,47 +352,15 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { this } - /** - * :: Experimental :: - * - * (Scala-specific) Sets the output of the streaming query to be processed using the provided - * function. This is supported only in the micro-batch execution modes (that is, when the - * trigger is not continuous). In every micro-batch, the provided function will be called in - * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. - * The batchId can be used to deduplicate and transactionally write the output - * (that is, the provided Dataset) to external systems. The output Dataset is guaranteed - * to be exactly the same for the same batchId (assuming all operations are deterministic - * in the query). - * - * @since 2.4.0 - */ + /** @inheritdoc */ @Evolving - def foreachBatch(function: (Dataset[T], Long) => Unit): DataStreamWriter[T] = { - this.source = SOURCE_NAME_FOREACH_BATCH + def foreachBatch(function: (Dataset[T], Long) => Unit): this.type = { + this.source = DataStreamWriter.SOURCE_NAME_FOREACH_BATCH if (function == null) throw new IllegalArgumentException("foreachBatch function cannot be null") this.foreachBatchWriter = function this } - /** - * :: Experimental :: - * - * (Java-specific) Sets the output of the streaming query to be processed using the provided - * function. This is supported only in the micro-batch execution modes (that is, when the - * trigger is not continuous). In every micro-batch, the provided function will be called in - * every micro-batch with (i) the output rows as a Dataset and (ii) the batch identifier. - * The batchId can be used to deduplicate and transactionally write the output - * (that is, the provided Dataset) to external systems. The output Dataset is guaranteed - * to be exactly the same for the same batchId (assuming all operations are deterministic - * in the query). - * - * @since 2.4.0 - */ - @Evolving - def foreachBatch(function: VoidFunction2[Dataset[T], java.lang.Long]): DataStreamWriter[T] = { - foreachBatch((batchDs: Dataset[T], batchId: Long) => function.call(batchDs, batchId)) - } - private def normalizedParCols: Option[Seq[String]] = partitioningColumns.map { cols => cols.map(normalize(_, "Partition")) } @@ -564,8 +375,8 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { * need to care about case sensitivity afterwards. */ private def normalize(columnName: String, columnType: String): String = { - val validColumnNames = df.logicalPlan.output.map(_.name) - validColumnNames.find(df.sparkSession.sessionState.analyzer.resolver(_, columnName)) + val validColumnNames = ds.logicalPlan.output.map(_.name) + validColumnNames.find(ds.sparkSession.sessionState.analyzer.resolver(_, columnName)) .getOrElse(throw QueryCompilationErrors.columnNotFoundInExistingColumnsError( columnType, columnName, validColumnNames)) } @@ -584,12 +395,28 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { } /////////////////////////////////////////////////////////////////////////////////////// - // Builder pattern config options + // Covariant Overrides /////////////////////////////////////////////////////////////////////////////////////// - private var source: String = df.sparkSession.sessionState.conf.defaultDataSourceName + /** @inheritdoc */ + override def option(key: String, value: Boolean): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Long): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Double): this.type = super.option(key, value) + + /** @inheritdoc */ + @Evolving + override def foreachBatch(function: VoidFunction2[Dataset[T], java.lang.Long]): this.type = + super.foreachBatch(function) + + /////////////////////////////////////////////////////////////////////////////////////// + // Builder pattern config options + /////////////////////////////////////////////////////////////////////////////////////// - private var tableName: String = null + private var source: String = ds.sparkSession.sessionState.conf.defaultDataSourceName private var outputMode: OutputMode = OutputMode.Append @@ -597,12 +424,12 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { private var extraOptions = CaseInsensitiveMap[String](Map.empty) - private var foreachWriter: ForeachWriter[Any] = null + private var foreachWriter: ForeachWriter[Any] = _ private var foreachWriterEncoder: ExpressionEncoder[Any] = ds.exprEnc.asInstanceOf[ExpressionEncoder[Any]] - private var foreachBatchWriter: (Dataset[T], Long) => Unit = null + private var foreachBatchWriter: (Dataset[T], Long) => Unit = _ private var partitioningColumns: Option[Seq[String]] = None @@ -610,14 +437,14 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) { } object DataStreamWriter { - val SOURCE_NAME_MEMORY = "memory" - val SOURCE_NAME_FOREACH = "foreach" - val SOURCE_NAME_FOREACH_BATCH = "foreachBatch" - val SOURCE_NAME_CONSOLE = "console" - val SOURCE_NAME_TABLE = "table" - val SOURCE_NAME_NOOP = "noop" + val SOURCE_NAME_MEMORY: String = "memory" + val SOURCE_NAME_FOREACH: String = "foreach" + val SOURCE_NAME_FOREACH_BATCH: String = "foreachBatch" + val SOURCE_NAME_CONSOLE: String = "console" + val SOURCE_NAME_TABLE: String = "table" + val SOURCE_NAME_NOOP: String = "noop" // these writer sources are also used for one-time query, hence allow temp checkpoint location - val SOURCES_ALLOW_ONE_TIME_QUERY = Seq(SOURCE_NAME_MEMORY, SOURCE_NAME_FOREACH, + val SOURCES_ALLOW_ONE_TIME_QUERY: Seq[String] = Seq(SOURCE_NAME_MEMORY, SOURCE_NAME_FOREACH, SOURCE_NAME_FOREACH_BATCH, SOURCE_NAME_CONSOLE, SOURCE_NAME_NOOP) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala index 3ab6d02f6b515..9d6fd2e28dea4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Evolving import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys.{CLASS_NAME, QUERY_ID, RUN_ID} -import org.apache.spark.sql.{DataFrame, SparkSession} +import org.apache.spark.sql.{Dataset, SparkSession} import org.apache.spark.sql.catalyst.catalog.CatalogTable import org.apache.spark.sql.catalyst.streaming.{WriteToStream, WriteToStreamStatement} import org.apache.spark.sql.connector.catalog.{Identifier, SupportsWrite, Table, TableCatalog} @@ -241,7 +241,7 @@ class StreamingQueryManager private[sql] ( private def createQuery( userSpecifiedName: Option[String], userSpecifiedCheckpointLocation: Option[String], - df: DataFrame, + df: Dataset[_], extraOptions: Map[String, String], sink: Table, outputMode: OutputMode, @@ -322,7 +322,7 @@ class StreamingQueryManager private[sql] ( private[sql] def startQuery( userSpecifiedName: Option[String], userSpecifiedCheckpointLocation: Option[String], - df: DataFrame, + df: Dataset[_], extraOptions: Map[String, String], sink: Table, outputMode: OutputMode, From 35e5d290deee9cf2a913571407e2257217e0e9e2 Mon Sep 17 00:00:00 2001 From: Chris Nauroth Date: Mon, 23 Sep 2024 21:35:32 -0700 Subject: [PATCH 113/189] [SPARK-49760][YARN] Correct handling of `SPARK_USER` env variable override in app master ### What changes were proposed in this pull request? This patch corrects handling of a user-supplied `SPARK_USER` environment variable in the YARN app master. Currently, the user-supplied value gets appended to the default, like a classpath entry. The patch fixes it by using only the user-supplied value. ### Why are the changes needed? Overriding the `SPARK_USER` environment variable in the YARN app master with configuration property `spark.yarn.appMasterEnv.SPARK_USER` currently results in an incorrect value. `Client#setupLaunchEnv` first sets a default in the environment map using the Hadoop user. After that, `YarnSparkHadoopUtil.addPathToEnvironment` sees the existing value in the map and interprets the user-supplied value as needing to be appended like a classpath entry. The end result is the Hadoop user appended with the classpath delimiter and user-supplied value, e.g. `cnauroth:overrideuser`. ### Does this PR introduce _any_ user-facing change? Yes, the app master now uses the user-supplied `SPARK_USER` if specified. (The default is still the Hadoop user.) ### How was this patch tested? * Existing unit tests pass. * Added new unit tests covering default and overridden `SPARK_USER` for the app master. The override test fails without this patch, and then passes after the patch is applied. * Manually tested in a live YARN cluster as shown below. Manual testing used the `DFSReadWriteTest` job with overrides of `SPARK_USER`: ``` spark-submit \ --deploy-mode cluster \ --files all-lines.txt \ --class org.apache.spark.examples.DFSReadWriteTest \ --conf spark.yarn.appMasterEnv.SPARK_USER=sparkuser_appMaster \ --conf spark.driverEnv.SPARK_USER=sparkuser_driver \ --conf spark.executorEnv.SPARK_USER=sparkuser_executor \ /usr/lib/spark/examples/jars/spark-examples.jar \ all-lines.txt /tmp/DFSReadWriteTest ``` Before the patch, we can see the app master's `SPARK_USER` mishandled by looking at the `_SUCCESS` file in HDFS: ``` hdfs dfs -ls -R /tmp/DFSReadWriteTest drwxr-xr-x - cnauroth:sparkuser_appMaster hadoop 0 2024-09-20 23:35 /tmp/DFSReadWriteTest/dfs_read_write_test -rw-r--r-- 1 cnauroth:sparkuser_appMaster hadoop 0 2024-09-20 23:35 /tmp/DFSReadWriteTest/dfs_read_write_test/_SUCCESS -rw-r--r-- 1 sparkuser_executor hadoop 2295080 2024-09-20 23:35 /tmp/DFSReadWriteTest/dfs_read_write_test/part-00000 -rw-r--r-- 1 sparkuser_executor hadoop 2288718 2024-09-20 23:35 /tmp/DFSReadWriteTest/dfs_read_write_test/part-00001 ``` After the patch, we can see it working correctly: ``` hdfs dfs -ls -R /tmp/DFSReadWriteTest drwxr-xr-x - sparkuser_appMaster hadoop 0 2024-09-23 17:13 /tmp/DFSReadWriteTest/dfs_read_write_test -rw-r--r-- 1 sparkuser_appMaster hadoop 0 2024-09-23 17:13 /tmp/DFSReadWriteTest/dfs_read_write_test/_SUCCESS -rw-r--r-- 1 sparkuser_executor hadoop 2295080 2024-09-23 17:13 /tmp/DFSReadWriteTest/dfs_read_write_test/part-00000 -rw-r--r-- 1 sparkuser_executor hadoop 2288718 2024-09-23 17:13 /tmp/DFSReadWriteTest/dfs_read_write_test/part-00001 ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48214 from cnauroth/SPARK-49760. Authored-by: Chris Nauroth Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/deploy/yarn/Client.scala | 7 +++++-- .../apache/spark/deploy/yarn/ClientSuite.scala | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala index b2c4d97bc7b07..8b621e82afe28 100644 --- a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala +++ b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala @@ -960,14 +960,13 @@ private[spark] class Client( /** * Set up the environment for launching our ApplicationMaster container. */ - private def setupLaunchEnv( + private[yarn] def setupLaunchEnv( stagingDirPath: Path, pySparkArchives: Seq[String]): HashMap[String, String] = { logInfo("Setting up the launch environment for our AM container") val env = new HashMap[String, String]() populateClasspath(args, hadoopConf, sparkConf, env, sparkConf.get(DRIVER_CLASS_PATH)) env("SPARK_YARN_STAGING_DIR") = stagingDirPath.toString - env("SPARK_USER") = UserGroupInformation.getCurrentUser().getShortUserName() env("SPARK_PREFER_IPV6") = Utils.preferIPv6.toString // Pick up any environment variables for the AM provided through spark.yarn.appMasterEnv.* @@ -977,6 +976,10 @@ private[spark] class Client( .map { case (k, v) => (k.substring(amEnvPrefix.length), v) } .foreach { case (k, v) => YarnSparkHadoopUtil.addPathToEnvironment(env, k, v) } + if (!env.contains("SPARK_USER")) { + env("SPARK_USER") = UserGroupInformation.getCurrentUser().getShortUserName() + } + // If pyFiles contains any .py files, we need to add LOCALIZED_PYTHON_DIR to the PYTHONPATH // of the container processes too. Add all non-.py files directly to PYTHONPATH. // diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/ClientSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/ClientSuite.scala index 78e84690900e1..93d6cc474d20f 100644 --- a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/ClientSuite.scala +++ b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/ClientSuite.scala @@ -29,6 +29,7 @@ import scala.jdk.CollectionConverters._ import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.{FileStatus, FileSystem, Path, PathFilter} import org.apache.hadoop.mapreduce.MRJobConfig +import org.apache.hadoop.security.UserGroupInformation import org.apache.hadoop.yarn.api.ApplicationConstants.Environment import org.apache.hadoop.yarn.api.protocolrecords.{GetNewApplicationResponse, SubmitApplicationRequest} import org.apache.hadoop.yarn.api.records._ @@ -739,6 +740,21 @@ class ClientSuite extends SparkFunSuite } } + test("SPARK-49760: default app master SPARK_USER") { + val sparkConf = new SparkConf() + val client = createClient(sparkConf) + val env = client.setupLaunchEnv(new Path("/staging/dir/path"), Seq()) + env("SPARK_USER") should be (UserGroupInformation.getCurrentUser().getShortUserName()) + } + + test("SPARK-49760: override app master SPARK_USER") { + val sparkConf = new SparkConf() + .set("spark.yarn.appMasterEnv.SPARK_USER", "overrideuser") + val client = createClient(sparkConf) + val env = client.setupLaunchEnv(new Path("/staging/dir/path"), Seq()) + env("SPARK_USER") should be ("overrideuser") + } + private val matching = Seq( ("files URI match test1", "file:///file1", "file:///file2"), ("files URI match test2", "file:///c:file1", "file://c:file2"), From 64ea50e87c70aea6b22a66ec1a0c98ae29a5dd81 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 24 Sep 2024 13:36:22 +0900 Subject: [PATCH 114/189] [SPARK-49607][PYTHON] Update the sampling approach for sampled based plots ### What changes were proposed in this pull request? 1, Update the sampling approach for sampled based plots 2, Eliminate "spark.sql.pyspark.plotting.sample_ratio" config ### Why are the changes needed? 1, to be consistent with the PS plotting; 2, the "spark.sql.pyspark.plotting.sample_ratio" config is not friendly to large scale data: the plotting backend cannot render large number of data points efficiently, and it is hard for users to set an appropriate sample ratio; ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48218 from zhengruifeng/py_plot_sampling. Authored-by: Ruifeng Zheng Signed-off-by: Hyukjin Kwon --- python/pyspark/sql/plot/core.py | 36 ++++++++++++++----- .../pyspark/sql/tests/plot/test_frame_plot.py | 14 +------- .../apache/spark/sql/internal/SQLConf.scala | 16 --------- 3 files changed, 28 insertions(+), 38 deletions(-) diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index ed22d02370ca6..eb00b8a04f977 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -50,27 +50,45 @@ def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": class PySparkSampledPlotBase: def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": - from pyspark.sql import SparkSession + from pyspark.sql import SparkSession, Observation, functions as F session = SparkSession.getActiveSession() if session is None: raise PySparkRuntimeError(errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) - sample_ratio = session.conf.get("spark.sql.pyspark.plotting.sample_ratio") max_rows = int( session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] ) - if sample_ratio is None: - fraction = 1 / (sdf.count() / max_rows) - fraction = min(1.0, fraction) - else: - fraction = float(sample_ratio) + observation = Observation("pyspark plotting") - sampled_sdf = sdf.sample(fraction=fraction) + rand_col_name = "__pyspark_plotting_sampled_plot_base_rand__" + id_col_name = "__pyspark_plotting_sampled_plot_base_id__" + + sampled_sdf = ( + sdf.observe(observation, F.count(F.lit(1)).alias("count")) + .select( + "*", + F.rand().alias(rand_col_name), + F.monotonically_increasing_id().alias(id_col_name), + ) + .sort(rand_col_name) + .limit(max_rows + 1) + .coalesce(1) + .sortWithinPartitions(id_col_name) + .drop(rand_col_name, id_col_name) + ) pdf = sampled_sdf.toPandas() - return pdf + if len(pdf) > max_rows: + try: + self.fraction = float(max_rows) / observation.get["count"] + except Exception: + pass + return pdf[:max_rows] + else: + self.fraction = 1.0 + return pdf class PySparkPlotAccessor: diff --git a/python/pyspark/sql/tests/plot/test_frame_plot.py b/python/pyspark/sql/tests/plot/test_frame_plot.py index f753b5ab3db72..2a6971e896292 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot.py @@ -39,23 +39,11 @@ def test_backend(self): ) def test_topn_max_rows(self): - try: + with self.sql_conf({"spark.sql.pyspark.plotting.max_rows": "1000"}): self.spark.conf.set("spark.sql.pyspark.plotting.max_rows", "1000") sdf = self.spark.range(2500) pdf = PySparkTopNPlotBase().get_top_n(sdf) self.assertEqual(len(pdf), 1000) - finally: - self.spark.conf.unset("spark.sql.pyspark.plotting.max_rows") - - def test_sampled_plot_with_ratio(self): - try: - self.spark.conf.set("spark.sql.pyspark.plotting.sample_ratio", "0.5") - data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2500)] - sdf = self.spark.createDataFrame(data) - pdf = PySparkSampledPlotBase().get_sampled(sdf) - self.assertEqual(round(len(pdf) / 2500, 1), 0.5) - finally: - self.spark.conf.unset("spark.sql.pyspark.plotting.sample_ratio") def test_sampled_plot_with_max_rows(self): data = [Row(a=i, b=i + 1, c=i + 2, d=i + 3) for i in range(2000)] diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 4d0930212b373..9d51afd064d10 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3178,20 +3178,6 @@ object SQLConf { .intConf .createWithDefault(1000) - val PYSPARK_PLOT_SAMPLE_RATIO = - buildConf("spark.sql.pyspark.plotting.sample_ratio") - .doc( - "The proportion of data that will be plotted for sample-based plots. It is determined " + - "based on spark.sql.pyspark.plotting.max_rows if not explicitly set." - ) - .version("4.0.0") - .doubleConf - .checkValue( - ratio => ratio >= 0.0 && ratio <= 1.0, - "The value should be between 0.0 and 1.0 inclusive." - ) - .createOptional - val ARROW_SPARKR_EXECUTION_ENABLED = buildConf("spark.sql.execution.arrow.sparkr.enabled") .doc("When true, make use of Apache Arrow for columnar data transfers in SparkR. " + @@ -5907,8 +5893,6 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS) - def pysparkPlotSampleRatio: Option[Double] = getConf(PYSPARK_PLOT_SAMPLE_RATIO) - def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED) def arrowPySparkFallbackEnabled: Boolean = getConf(ARROW_PYSPARK_FALLBACK_ENABLED) From 438a6e7782ece23492928cfbb2d01e14104dfd9a Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 23 Sep 2024 21:39:27 -0700 Subject: [PATCH 115/189] [SPARK-49753][BUILD] Upgrade ZSTD-JNI to 1.5.6-6 ### What changes were proposed in this pull request? The pr aims to upgrade `zstd-jni` from `1.5.6-5` to `1.5.6-6`. ### Why are the changes needed? The new version allow including compression level when training a dictionary: https://github.com/luben/zstd-jni/commit/3ca26eed6c84fb09c382854ead527188e643e206#diff-bd5c0f62db7cb85cac88c7b6cfad1c0e5e2f433ba45097761654829627b7a31c All changes in the new version are as follows: - https://github.com/luben/zstd-jni/compare/v1.5.6-5...v1.5.6-6 ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GitHub Actions ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48204 from LuciferYang/zstd-jni-1.5.6-6. Lead-authored-by: yangjie01 Co-authored-by: YangJie Signed-off-by: Dongjoon Hyun --- .../ZStandardBenchmark-jdk21-results.txt | 56 +++++++++---------- .../benchmarks/ZStandardBenchmark-results.txt | 56 +++++++++---------- dev/deps/spark-deps-hadoop-3-hive-2.3 | 2 +- pom.xml | 2 +- 4 files changed, 58 insertions(+), 58 deletions(-) diff --git a/core/benchmarks/ZStandardBenchmark-jdk21-results.txt b/core/benchmarks/ZStandardBenchmark-jdk21-results.txt index b3bffea826e5f..f6bd681451d5e 100644 --- a/core/benchmarks/ZStandardBenchmark-jdk21-results.txt +++ b/core/benchmarks/ZStandardBenchmark-jdk21-results.txt @@ -2,48 +2,48 @@ Benchmark ZStandardCompressionCodec ================================================================================================ -OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Benchmark ZStandardCompressionCodec: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative -------------------------------------------------------------------------------------------------------------------------------------- -Compression 10000 times at level 1 without buffer pool 657 670 14 0.0 65699.2 1.0X -Compression 10000 times at level 2 without buffer pool 697 697 1 0.0 69673.4 0.9X -Compression 10000 times at level 3 without buffer pool 799 802 3 0.0 79855.2 0.8X -Compression 10000 times at level 1 with buffer pool 593 595 1 0.0 59326.9 1.1X -Compression 10000 times at level 2 with buffer pool 622 624 3 0.0 62194.1 1.1X -Compression 10000 times at level 3 with buffer pool 732 733 1 0.0 73178.6 0.9X +Compression 10000 times at level 1 without buffer pool 659 676 16 0.0 65860.7 1.0X +Compression 10000 times at level 2 without buffer pool 721 723 2 0.0 72135.5 0.9X +Compression 10000 times at level 3 without buffer pool 815 816 1 0.0 81500.6 0.8X +Compression 10000 times at level 1 with buffer pool 608 609 0 0.0 60846.6 1.1X +Compression 10000 times at level 2 with buffer pool 645 647 3 0.0 64476.3 1.0X +Compression 10000 times at level 3 with buffer pool 746 746 1 0.0 74584.0 0.9X -OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Benchmark ZStandardCompressionCodec: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------------------------ -Decompression 10000 times from level 1 without buffer pool 813 820 11 0.0 81273.2 1.0X -Decompression 10000 times from level 2 without buffer pool 810 813 3 0.0 80986.2 1.0X -Decompression 10000 times from level 3 without buffer pool 812 813 2 0.0 81183.1 1.0X -Decompression 10000 times from level 1 with buffer pool 746 747 2 0.0 74568.7 1.1X -Decompression 10000 times from level 2 with buffer pool 744 746 2 0.0 74414.5 1.1X -Decompression 10000 times from level 3 with buffer pool 745 746 1 0.0 74538.6 1.1X +Decompression 10000 times from level 1 without buffer pool 828 829 1 0.0 82822.6 1.0X +Decompression 10000 times from level 2 without buffer pool 829 829 1 0.0 82900.7 1.0X +Decompression 10000 times from level 3 without buffer pool 828 833 8 0.0 82784.4 1.0X +Decompression 10000 times from level 1 with buffer pool 758 760 2 0.0 75756.5 1.1X +Decompression 10000 times from level 2 with buffer pool 758 758 1 0.0 75772.3 1.1X +Decompression 10000 times from level 3 with buffer pool 759 759 0 0.0 75852.7 1.1X -OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Parallel Compression at level 3: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -Parallel Compression with 0 workers 48 49 1 0.0 374256.1 1.0X -Parallel Compression with 1 workers 34 36 3 0.0 267557.3 1.4X -Parallel Compression with 2 workers 34 38 2 0.0 263684.3 1.4X -Parallel Compression with 4 workers 37 39 2 0.0 289956.1 1.3X -Parallel Compression with 8 workers 39 41 1 0.0 306975.2 1.2X -Parallel Compression with 16 workers 44 45 1 0.0 340992.0 1.1X +Parallel Compression with 0 workers 58 59 1 0.0 452489.9 1.0X +Parallel Compression with 1 workers 42 45 4 0.0 330066.0 1.4X +Parallel Compression with 2 workers 40 42 1 0.0 312560.3 1.4X +Parallel Compression with 4 workers 40 42 2 0.0 308802.7 1.5X +Parallel Compression with 8 workers 41 45 3 0.0 321331.3 1.4X +Parallel Compression with 16 workers 44 45 1 0.0 343311.5 1.3X -OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 21.0.4+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Parallel Compression at level 9: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -Parallel Compression with 0 workers 156 158 1 0.0 1220760.5 1.0X -Parallel Compression with 1 workers 191 192 2 0.0 1495168.2 0.8X -Parallel Compression with 2 workers 111 117 5 0.0 864459.9 1.4X -Parallel Compression with 4 workers 106 109 2 0.0 831025.5 1.5X -Parallel Compression with 8 workers 112 115 2 0.0 875732.7 1.4X -Parallel Compression with 16 workers 110 114 2 0.0 858160.9 1.4X +Parallel Compression with 0 workers 158 160 2 0.0 1234257.6 1.0X +Parallel Compression with 1 workers 193 194 1 0.0 1507686.4 0.8X +Parallel Compression with 2 workers 113 127 11 0.0 881068.0 1.4X +Parallel Compression with 4 workers 109 111 2 0.0 849241.3 1.5X +Parallel Compression with 8 workers 111 115 3 0.0 869455.2 1.4X +Parallel Compression with 16 workers 113 116 2 0.0 881832.5 1.4X diff --git a/core/benchmarks/ZStandardBenchmark-results.txt b/core/benchmarks/ZStandardBenchmark-results.txt index b230f825fecac..136f0333590cc 100644 --- a/core/benchmarks/ZStandardBenchmark-results.txt +++ b/core/benchmarks/ZStandardBenchmark-results.txt @@ -2,48 +2,48 @@ Benchmark ZStandardCompressionCodec ================================================================================================ -OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Benchmark ZStandardCompressionCodec: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative -------------------------------------------------------------------------------------------------------------------------------------- -Compression 10000 times at level 1 without buffer pool 638 638 0 0.0 63765.0 1.0X -Compression 10000 times at level 2 without buffer pool 675 676 1 0.0 67529.4 0.9X -Compression 10000 times at level 3 without buffer pool 775 783 11 0.0 77531.6 0.8X -Compression 10000 times at level 1 with buffer pool 572 573 1 0.0 57223.2 1.1X -Compression 10000 times at level 2 with buffer pool 603 605 1 0.0 60323.7 1.1X -Compression 10000 times at level 3 with buffer pool 720 727 6 0.0 71980.9 0.9X +Compression 10000 times at level 1 without buffer pool 257 259 2 0.0 25704.2 1.0X +Compression 10000 times at level 2 without buffer pool 674 676 2 0.0 67396.3 0.4X +Compression 10000 times at level 3 without buffer pool 775 787 11 0.0 77497.9 0.3X +Compression 10000 times at level 1 with buffer pool 573 574 0 0.0 57347.3 0.4X +Compression 10000 times at level 2 with buffer pool 602 603 2 0.0 60162.8 0.4X +Compression 10000 times at level 3 with buffer pool 722 725 3 0.0 72247.3 0.4X -OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Benchmark ZStandardCompressionCodec: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------------------------ -Decompression 10000 times from level 1 without buffer pool 584 585 1 0.0 58381.0 1.0X -Decompression 10000 times from level 2 without buffer pool 585 585 0 0.0 58465.9 1.0X -Decompression 10000 times from level 3 without buffer pool 585 586 1 0.0 58499.5 1.0X -Decompression 10000 times from level 1 with buffer pool 534 534 0 0.0 53375.7 1.1X -Decompression 10000 times from level 2 with buffer pool 533 533 0 0.0 53312.3 1.1X -Decompression 10000 times from level 3 with buffer pool 533 533 1 0.0 53255.1 1.1X +Decompression 10000 times from level 1 without buffer pool 176 177 1 0.1 17641.2 1.0X +Decompression 10000 times from level 2 without buffer pool 176 178 1 0.1 17628.9 1.0X +Decompression 10000 times from level 3 without buffer pool 175 176 0 0.1 17506.1 1.0X +Decompression 10000 times from level 1 with buffer pool 151 152 1 0.1 15051.5 1.2X +Decompression 10000 times from level 2 with buffer pool 150 151 1 0.1 14998.0 1.2X +Decompression 10000 times from level 3 with buffer pool 150 151 0 0.1 15019.4 1.2X -OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Parallel Compression at level 3: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -Parallel Compression with 0 workers 46 48 1 0.0 360483.5 1.0X -Parallel Compression with 1 workers 34 36 2 0.0 265816.1 1.4X -Parallel Compression with 2 workers 33 36 2 0.0 254525.8 1.4X -Parallel Compression with 4 workers 34 37 1 0.0 266270.8 1.4X -Parallel Compression with 8 workers 37 39 1 0.0 289289.2 1.2X -Parallel Compression with 16 workers 41 43 1 0.0 320243.3 1.1X +Parallel Compression with 0 workers 57 57 0 0.0 444425.2 1.0X +Parallel Compression with 1 workers 42 44 3 0.0 325107.6 1.4X +Parallel Compression with 2 workers 38 39 2 0.0 294840.0 1.5X +Parallel Compression with 4 workers 36 37 1 0.0 282143.1 1.6X +Parallel Compression with 8 workers 39 40 1 0.0 303793.6 1.5X +Parallel Compression with 16 workers 41 43 1 0.0 324165.5 1.4X -OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.5.0-1025-azure +OpenJDK 64-Bit Server VM 17.0.12+7-LTS on Linux 6.8.0-1014-azure AMD EPYC 7763 64-Core Processor Parallel Compression at level 9: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -Parallel Compression with 0 workers 154 156 2 0.0 1205934.0 1.0X -Parallel Compression with 1 workers 191 194 4 0.0 1495729.9 0.8X -Parallel Compression with 2 workers 110 114 5 0.0 859158.9 1.4X -Parallel Compression with 4 workers 105 108 3 0.0 822932.2 1.5X -Parallel Compression with 8 workers 109 113 2 0.0 851560.0 1.4X -Parallel Compression with 16 workers 111 115 2 0.0 870695.9 1.4X +Parallel Compression with 0 workers 156 158 1 0.0 1220298.8 1.0X +Parallel Compression with 1 workers 188 189 1 0.0 1467911.4 0.8X +Parallel Compression with 2 workers 111 118 7 0.0 866985.2 1.4X +Parallel Compression with 4 workers 106 109 2 0.0 827592.1 1.5X +Parallel Compression with 8 workers 114 116 2 0.0 888419.5 1.4X +Parallel Compression with 16 workers 111 115 2 0.0 868463.5 1.4X diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 419625f48fa11..88526995293f5 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -280,4 +280,4 @@ xz/1.10//xz-1.10.jar zjsonpatch/0.3.0//zjsonpatch-0.3.0.jar zookeeper-jute/3.9.2//zookeeper-jute-3.9.2.jar zookeeper/3.9.2//zookeeper-3.9.2.jar -zstd-jni/1.5.6-5//zstd-jni-1.5.6-5.jar +zstd-jni/1.5.6-6//zstd-jni-1.5.6-6.jar diff --git a/pom.xml b/pom.xml index b7c87beec0f92..131e754da8157 100644 --- a/pom.xml +++ b/pom.xml @@ -835,7 +835,7 @@ com.github.luben zstd-jni - 1.5.6-5 + 1.5.6-6 com.clearspring.analytics From 6bdd151d57759d73870f20780fc54ab2aa250409 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Tue, 24 Sep 2024 15:40:38 +0800 Subject: [PATCH 116/189] [SPARK-49694][PYTHON][CONNECT] Support scatter plots ### What changes were proposed in this pull request? Support scatter plots with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. Scatter plots are supported as shown below. ```py >>> data = [(5.1, 3.5, 0), (4.9, 3.0, 0), (7.0, 3.2, 1), (6.4, 3.2, 1), (5.9, 3.0, 2)] >>> columns = ["length", "width", "species"] >>> sdf = spark.createDataFrame(data, columns) >>> fig = sdf.plot(kind="scatter", x="length", y="width") # or fig = sdf.plot.scatter(x="length", y="width") >>> fig.show() ``` ![newplot (6)](https://github.com/user-attachments/assets/deef452b-74d1-4f6d-b1ae-60722f3c2b17) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48219 from xinrong-meng/plot_scatter. Authored-by: Xinrong Meng Signed-off-by: Xinrong Meng --- python/pyspark/sql/plot/core.py | 34 +++++++++++++++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 19 +++++++++++ 2 files changed, 53 insertions(+) diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index eb00b8a04f977..0a3a0101e1898 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -96,6 +96,7 @@ class PySparkPlotAccessor: "bar": PySparkTopNPlotBase().get_top_n, "barh": PySparkTopNPlotBase().get_top_n, "line": PySparkSampledPlotBase().get_sampled, + "scatter": PySparkSampledPlotBase().get_sampled, } _backends = {} # type: ignore[var-annotated] @@ -230,3 +231,36 @@ def barh(self, x: str, y: Union[str, list[str]], **kwargs: Any) -> "Figure": ... ) # doctest: +SKIP """ return self(kind="barh", x=x, y=y, **kwargs) + + def scatter(self, x: str, y: str, **kwargs: Any) -> "Figure": + """ + Create a scatter plot with varying marker point size and color. + + The coordinates of each point are defined by two dataframe columns and + filled circles are used to represent each point. This kind of plot is + useful to see complex correlations between two variables. Points could + be for instance natural 2D coordinates like longitude and latitude in + a map or, in general, any pair of metrics that can be plotted against + each other. + + Parameters + ---------- + x : str + Name of column to use as horizontal coordinates for each point. + y : str or list of str + Name of column to use as vertical coordinates for each point. + **kwargs: Optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + >>> data = [(5.1, 3.5, 0), (4.9, 3.0, 0), (7.0, 3.2, 1), (6.4, 3.2, 1), (5.9, 3.0, 2)] + >>> columns = ['length', 'width', 'species'] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.scatter(x='length', y='width') # doctest: +SKIP + """ + return self(kind="scatter", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py index 1c52c93a23d3a..ccfe1a75424e0 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -28,6 +28,12 @@ def sdf(self): columns = ["category", "int_val", "float_val"] return self.spark.createDataFrame(data, columns) + @property + def sdf2(self): + data = [(5.1, 3.5, 0), (4.9, 3.0, 0), (7.0, 3.2, 1), (6.4, 3.2, 1), (5.9, 3.0, 2)] + columns = ["length", "width", "species"] + return self.spark.createDataFrame(data, columns) + def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name=""): if kind == "line": self.assertEqual(fig_data["mode"], "lines") @@ -37,6 +43,9 @@ def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name= elif kind == "barh": self.assertEqual(fig_data["type"], "bar") self.assertEqual(fig_data["orientation"], "h") + elif kind == "scatter": + self.assertEqual(fig_data["type"], "scatter") + self.assertEqual(fig_data["orientation"], "v") self.assertEqual(fig_data["xaxis"], "x") self.assertEqual(list(fig_data["x"]), expected_x) @@ -79,6 +88,16 @@ def test_barh_plot(self): self._check_fig_data("barh", fig["data"][0], [10, 30, 20], ["A", "B", "C"], "int_val") self._check_fig_data("barh", fig["data"][1], [1.5, 2.5, 3.5], ["A", "B", "C"], "float_val") + def test_scatter_plot(self): + fig = self.sdf2.plot(kind="scatter", x="length", y="width") + self._check_fig_data( + "scatter", fig["data"][0], [5.1, 4.9, 7.0, 6.4, 5.9], [3.5, 3.0, 3.2, 3.2, 3.0] + ) + fig = self.sdf2.plot.scatter(x="width", y="length") + self._check_fig_data( + "scatter", fig["data"][0], [3.5, 3.0, 3.2, 3.2, 3.0], [5.1, 4.9, 7.0, 6.4, 5.9] + ) + class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): pass From 982028ea7fc61d7aa84756aa46860ebb49bfe9d1 Mon Sep 17 00:00:00 2001 From: Haejoon Lee Date: Tue, 24 Sep 2024 17:21:59 +0900 Subject: [PATCH 117/189] [SPARK-49609][PYTHON][CONNECT] Add API compatibility check between Classic and Connect ### What changes were proposed in this pull request? This PR proposes to add API compatibility check between Classic and Connect. This PR also includes updating both APIs to the same signature. ### Why are the changes needed? APIs supported on both Spark Connect and Spark Classic should guarantee the same signature, such as argument and return types. For example, test would fail when the signature of API is mismatched: ``` Signature mismatch in Column method 'dropFields' Classic: (self, *fieldNames: str) -> pyspark.sql.column.Column Connect: (self, *fieldNames: 'ColumnOrName') -> pyspark.sql.column.Column pyspark.sql.column.Column> != pyspark.sql.column.Column> Expected : pyspark.sql.column.Column> Actual : pyspark.sql.column.Column> ``` ### Does this PR introduce _any_ user-facing change? No, it is a test to prevent future API behavior inconsistencies between Classic and Connect. ### How was this patch tested? Added UTs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48085 from itholic/SPARK-49609. Authored-by: Haejoon Lee Signed-off-by: Haejoon Lee --- dev/sparktestsupport/modules.py | 1 + python/pyspark/sql/classic/dataframe.py | 6 +- python/pyspark/sql/connect/dataframe.py | 26 ++- python/pyspark/sql/session.py | 3 +- .../sql/tests/test_connect_compatibility.py | 188 ++++++++++++++++++ 5 files changed, 209 insertions(+), 15 deletions(-) create mode 100644 python/pyspark/sql/tests/test_connect_compatibility.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index b9a4bed715f67..eda6b063350e5 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -550,6 +550,7 @@ def __hash__(self): "pyspark.sql.tests.test_resources", "pyspark.sql.tests.plot.test_frame_plot", "pyspark.sql.tests.plot.test_frame_plot_plotly", + "pyspark.sql.tests.test_connect_compatibility", ], ) diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index a2778cbc32c4c..23484fcf0051f 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -1068,7 +1068,7 @@ def selectExpr(self, *expr: Union[str, List[str]]) -> ParentDataFrame: jdf = self._jdf.selectExpr(self._jseq(expr)) return DataFrame(jdf, self.sparkSession) - def filter(self, condition: "ColumnOrName") -> ParentDataFrame: + def filter(self, condition: Union[Column, str]) -> ParentDataFrame: if isinstance(condition, str): jdf = self._jdf.filter(condition) elif isinstance(condition, Column): @@ -1809,10 +1809,10 @@ def groupby(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": # type: ign def drop_duplicates(self, subset: Optional[List[str]] = None) -> ParentDataFrame: return self.dropDuplicates(subset) - def writeTo(self, table: str) -> DataFrameWriterV2: + def writeTo(self, table: str) -> "DataFrameWriterV2": return DataFrameWriterV2(self, table) - def mergeInto(self, table: str, condition: Column) -> MergeIntoWriter: + def mergeInto(self, table: str, condition: Column) -> "MergeIntoWriter": return MergeIntoWriter(self, table, condition) def pandas_api( diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 59d79decf6690..cb37af8868aad 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -535,7 +535,7 @@ def groupby(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": def groupby(self, __cols: Union[List[Column], List[str], List[int]]) -> "GroupedData": ... - def groupBy(self, *cols: "ColumnOrNameOrOrdinal") -> GroupedData: + def groupBy(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": if len(cols) == 1 and isinstance(cols[0], list): cols = cols[0] @@ -570,7 +570,7 @@ def rollup(self, *cols: "ColumnOrName") -> "GroupedData": def rollup(self, __cols: Union[List[Column], List[str]]) -> "GroupedData": ... - def rollup(self, *cols: "ColumnOrName") -> "GroupedData": # type: ignore[misc] + def rollup(self, *cols: "ColumnOrNameOrOrdinal") -> "GroupedData": # type: ignore[misc] _cols: List[Column] = [] for c in cols: if isinstance(c, Column): @@ -731,8 +731,8 @@ def _convert_col(df: ParentDataFrame, col: "ColumnOrName") -> Column: session=self._session, ) - def limit(self, n: int) -> ParentDataFrame: - res = DataFrame(plan.Limit(child=self._plan, limit=n), session=self._session) + def limit(self, num: int) -> ParentDataFrame: + res = DataFrame(plan.Limit(child=self._plan, limit=num), session=self._session) res._cached_schema = self._cached_schema return res @@ -931,7 +931,11 @@ def _show_string( )._to_table() return table[0][0].as_py() - def withColumns(self, colsMap: Dict[str, Column]) -> ParentDataFrame: + def withColumns(self, *colsMap: Dict[str, Column]) -> ParentDataFrame: + # Below code is to help enable kwargs in future. + assert len(colsMap) == 1 + colsMap = colsMap[0] # type: ignore[assignment] + if not isinstance(colsMap, dict): raise PySparkTypeError( errorClass="NOT_DICT", @@ -1256,7 +1260,7 @@ def intersectAll(self, other: ParentDataFrame) -> ParentDataFrame: res._cached_schema = self._merge_cached_schema(other) return res - def where(self, condition: Union[Column, str]) -> ParentDataFrame: + def where(self, condition: "ColumnOrName") -> ParentDataFrame: if not isinstance(condition, (str, Column)): raise PySparkTypeError( errorClass="NOT_COLUMN_OR_STR", @@ -2193,7 +2197,7 @@ def cb(ei: "ExecutionInfo") -> None: return DataFrameWriterV2(self._plan, self._session, table, cb) - def mergeInto(self, table: str, condition: Column) -> MergeIntoWriter: + def mergeInto(self, table: str, condition: Column) -> "MergeIntoWriter": def cb(ei: "ExecutionInfo") -> None: self._execution_info = ei @@ -2201,10 +2205,10 @@ def cb(ei: "ExecutionInfo") -> None: self._plan, self._session, table, condition, cb # type: ignore[arg-type] ) - def offset(self, n: int) -> ParentDataFrame: - return DataFrame(plan.Offset(child=self._plan, offset=n), session=self._session) + def offset(self, num: int) -> ParentDataFrame: + return DataFrame(plan.Offset(child=self._plan, offset=num), session=self._session) - def checkpoint(self, eager: bool = True) -> "DataFrame": + def checkpoint(self, eager: bool = True) -> ParentDataFrame: cmd = plan.Checkpoint(child=self._plan, local=False, eager=eager) _, properties, self._execution_info = self._session.client.execute_command( cmd.command(self._session.client) @@ -2214,7 +2218,7 @@ def checkpoint(self, eager: bool = True) -> "DataFrame": assert isinstance(checkpointed._plan, plan.CachedRemoteRelation) return checkpointed - def localCheckpoint(self, eager: bool = True) -> "DataFrame": + def localCheckpoint(self, eager: bool = True) -> ParentDataFrame: cmd = plan.Checkpoint(child=self._plan, local=True, eager=eager) _, properties, self._execution_info = self._session.client.execute_command( cmd.command(self._session.client) diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index b513d8d4111b9..96344efba2d2a 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -77,6 +77,7 @@ from pyspark.sql.udf import UDFRegistration from pyspark.sql.udtf import UDTFRegistration from pyspark.sql.datasource import DataSourceRegistration + from pyspark.sql.dataframe import DataFrame as ParentDataFrame # Running MyPy type checks will always require pandas and # other dependencies so importing here is fine. @@ -1641,7 +1642,7 @@ def prepare(obj: Any) -> Any: def sql( self, sqlQuery: str, args: Optional[Union[Dict[str, Any], List]] = None, **kwargs: Any - ) -> DataFrame: + ) -> "ParentDataFrame": """Returns a :class:`DataFrame` representing the result of the given query. When ``kwargs`` is specified, this method formats the given string by using the Python standard formatter. The method binds named parameters to SQL literals or diff --git a/python/pyspark/sql/tests/test_connect_compatibility.py b/python/pyspark/sql/tests/test_connect_compatibility.py new file mode 100644 index 0000000000000..ca1f828ef4d78 --- /dev/null +++ b/python/pyspark/sql/tests/test_connect_compatibility.py @@ -0,0 +1,188 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +import inspect + +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.sql.classic.dataframe import DataFrame as ClassicDataFrame +from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame +from pyspark.sql.classic.column import Column as ClassicColumn +from pyspark.sql.connect.column import Column as ConnectColumn +from pyspark.sql.session import SparkSession as ClassicSparkSession +from pyspark.sql.connect.session import SparkSession as ConnectSparkSession + + +class ConnectCompatibilityTestsMixin: + def get_public_methods(self, cls): + """Get public methods of a class.""" + return { + name: method + for name, method in inspect.getmembers(cls, predicate=inspect.isfunction) + if not name.startswith("_") + } + + def get_public_properties(self, cls): + """Get public properties of a class.""" + return { + name: member + for name, member in inspect.getmembers(cls) + if isinstance(member, property) and not name.startswith("_") + } + + def test_signature_comparison_between_classic_and_connect(self): + def compare_method_signatures(classic_cls, connect_cls, cls_name): + """Compare method signatures between classic and connect classes.""" + classic_methods = self.get_public_methods(classic_cls) + connect_methods = self.get_public_methods(connect_cls) + + common_methods = set(classic_methods.keys()) & set(connect_methods.keys()) + + for method in common_methods: + classic_signature = inspect.signature(classic_methods[method]) + connect_signature = inspect.signature(connect_methods[method]) + + # createDataFrame cannot be the same since RDD is not supported from Spark Connect + if not method == "createDataFrame": + self.assertEqual( + classic_signature, + connect_signature, + f"Signature mismatch in {cls_name} method '{method}'\n" + f"Classic: {classic_signature}\n" + f"Connect: {connect_signature}", + ) + + # DataFrame API signature comparison + compare_method_signatures(ClassicDataFrame, ConnectDataFrame, "DataFrame") + + # Column API signature comparison + compare_method_signatures(ClassicColumn, ConnectColumn, "Column") + + # SparkSession API signature comparison + compare_method_signatures(ClassicSparkSession, ConnectSparkSession, "SparkSession") + + def test_property_comparison_between_classic_and_connect(self): + def compare_property_lists(classic_cls, connect_cls, cls_name, expected_missing_properties): + """Compare properties between classic and connect classes.""" + classic_properties = self.get_public_properties(classic_cls) + connect_properties = self.get_public_properties(connect_cls) + + # Identify missing properties + classic_only_properties = set(classic_properties.keys()) - set( + connect_properties.keys() + ) + + # Compare the actual missing properties with the expected ones + self.assertEqual( + classic_only_properties, + expected_missing_properties, + f"{cls_name}: Unexpected missing properties in Connect: {classic_only_properties}", + ) + + # Expected missing properties for DataFrame + expected_missing_properties_for_dataframe = {"sql_ctx", "isStreaming"} + + # DataFrame properties comparison + compare_property_lists( + ClassicDataFrame, + ConnectDataFrame, + "DataFrame", + expected_missing_properties_for_dataframe, + ) + + # Expected missing properties for Column (if any, replace with actual values) + expected_missing_properties_for_column = set() + + # Column properties comparison + compare_property_lists( + ClassicColumn, ConnectColumn, "Column", expected_missing_properties_for_column + ) + + # Expected missing properties for SparkSession + expected_missing_properties_for_spark_session = {"sparkContext", "version"} + + # SparkSession properties comparison + compare_property_lists( + ClassicSparkSession, + ConnectSparkSession, + "SparkSession", + expected_missing_properties_for_spark_session, + ) + + def test_missing_methods(self): + def check_missing_methods(classic_cls, connect_cls, cls_name, expected_missing_methods): + """Check for expected missing methods between classic and connect classes.""" + classic_methods = self.get_public_methods(classic_cls) + connect_methods = self.get_public_methods(connect_cls) + + # Identify missing methods + classic_only_methods = set(classic_methods.keys()) - set(connect_methods.keys()) + + # Compare the actual missing methods with the expected ones + self.assertEqual( + classic_only_methods, + expected_missing_methods, + f"{cls_name}: Unexpected missing methods in Connect: {classic_only_methods}", + ) + + # Expected missing methods for DataFrame + expected_missing_methods_for_dataframe = { + "inputFiles", + "isLocal", + "semanticHash", + "isEmpty", + } + + # DataFrame missing method check + check_missing_methods( + ClassicDataFrame, ConnectDataFrame, "DataFrame", expected_missing_methods_for_dataframe + ) + + # Expected missing methods for Column (if any, replace with actual values) + expected_missing_methods_for_column = set() + + # Column missing method check + check_missing_methods( + ClassicColumn, ConnectColumn, "Column", expected_missing_methods_for_column + ) + + # Expected missing methods for SparkSession (if any, replace with actual values) + expected_missing_methods_for_spark_session = {"newSession"} + + # SparkSession missing method check + check_missing_methods( + ClassicSparkSession, + ConnectSparkSession, + "SparkSession", + expected_missing_methods_for_spark_session, + ) + + +class ConnectCompatibilityTests(ConnectCompatibilityTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.test_connect_compatibility import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) From 73d6bd7c35b599690d40efe306eea0774f272ba8 Mon Sep 17 00:00:00 2001 From: Anish Shrigondekar Date: Tue, 24 Sep 2024 19:06:12 +0900 Subject: [PATCH 118/189] [SPARK-49630][SS] Add flatten option to process collection types with state data source reader ### What changes were proposed in this pull request? Add flatten option to process collection types with state data source reader ### Why are the changes needed? Changes are needed to process entries row-by-row in case we don't have enough memory to fit these collections inside a single row ### Does this PR introduce _any_ user-facing change? Yes Users can provide the following query option: ``` val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, ) .option(StateSourceOptions.STATE_VAR_NAME, ) .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, ) .load() ``` ### How was this patch tested? Added unit tests ``` [info] Run completed in 1 minute, 10 seconds. [info] Total number of tests run: 12 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 12, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48110 from anishshri-db/task/SPARK-49630. Authored-by: Anish Shrigondekar Signed-off-by: Jungtaek Lim --- .../v2/state/StateDataSource.scala | 19 +- .../v2/state/StatePartitionReader.scala | 55 +--- .../v2/state/utils/SchemaUtil.scala | 264 ++++++++++++------ .../v2/state/StateDataSourceReadSuite.scala | 19 ++ ...ateDataSourceTransformWithStateSuite.scala | 88 +++++- .../streaming/TransformWithStateSuite.scala | 8 +- 6 files changed, 320 insertions(+), 133 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala index 50b90641d309b..429464ea5438d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala @@ -303,13 +303,15 @@ case class StateSourceOptions( readChangeFeed: Boolean, fromSnapshotOptions: Option[FromSnapshotOptions], readChangeFeedOptions: Option[ReadChangeFeedOptions], - stateVarName: Option[String]) { + stateVarName: Option[String], + flattenCollectionTypes: Boolean) { def stateCheckpointLocation: Path = new Path(resolvedCpLocation, DIR_NAME_STATE) override def toString: String = { var desc = s"StateSourceOptions(checkpointLocation=$resolvedCpLocation, batchId=$batchId, " + s"operatorId=$operatorId, storeName=$storeName, joinSide=$joinSide, " + - s"stateVarName=${stateVarName.getOrElse("None")}" + s"stateVarName=${stateVarName.getOrElse("None")}, +" + + s"flattenCollectionTypes=$flattenCollectionTypes" if (fromSnapshotOptions.isDefined) { desc += s", snapshotStartBatchId=${fromSnapshotOptions.get.snapshotStartBatchId}" desc += s", snapshotPartitionId=${fromSnapshotOptions.get.snapshotPartitionId}" @@ -334,6 +336,7 @@ object StateSourceOptions extends DataSourceOptions { val CHANGE_START_BATCH_ID = newOption("changeStartBatchId") val CHANGE_END_BATCH_ID = newOption("changeEndBatchId") val STATE_VAR_NAME = newOption("stateVarName") + val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes") object JoinSideValues extends Enumeration { type JoinSideValues = Value @@ -374,6 +377,15 @@ object StateSourceOptions extends DataSourceOptions { val stateVarName = Option(options.get(STATE_VAR_NAME)) .map(_.trim) + val flattenCollectionTypes = try { + Option(options.get(FLATTEN_COLLECTION_TYPES)) + .map(_.toBoolean).getOrElse(true) + } catch { + case _: IllegalArgumentException => + throw StateDataSourceErrors.invalidOptionValue(FLATTEN_COLLECTION_TYPES, + "Boolean value is expected") + } + val joinSide = try { Option(options.get(JOIN_SIDE)) .map(JoinSideValues.withName).getOrElse(JoinSideValues.none) @@ -477,7 +489,8 @@ object StateSourceOptions extends DataSourceOptions { StateSourceOptions( resolvedCpLocation, batchId.get, operatorId, storeName, joinSide, - readChangeFeed, fromSnapshotOptions, readChangeFeedOptions, stateVarName) + readChangeFeed, fromSnapshotOptions, readChangeFeedOptions, stateVarName, + flattenCollectionTypes) } private def resolvedCheckpointLocation( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala index 24166a46bbd39..ae12b18c1f627 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala @@ -19,7 +19,6 @@ package org.apache.spark.sql.execution.datasources.v2.state import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow} -import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory} import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil import org.apache.spark.sql.execution.streaming.{StateVariableType, TransformWithStateVariableInfo} @@ -75,9 +74,11 @@ abstract class StatePartitionReaderBase( StructType(Array(StructField("__dummy__", NullType))) protected val keySchema = { - if (!SchemaUtil.isMapStateVariable(stateVariableInfoOpt)) { + if (SchemaUtil.checkVariableType(stateVariableInfoOpt, StateVariableType.MapState)) { + SchemaUtil.getCompositeKeySchema(schema, partition.sourceOptions) + } else { SchemaUtil.getSchemaAsDataType(schema, "key").asInstanceOf[StructType] - } else SchemaUtil.getCompositeKeySchema(schema) + } } protected val valueSchema = if (stateVariableInfoOpt.isDefined) { @@ -98,12 +99,8 @@ abstract class StatePartitionReaderBase( false } - val useMultipleValuesPerKey = if (stateVariableInfoOpt.isDefined && - stateVariableInfoOpt.get.stateVariableType == StateVariableType.ListState) { - true - } else { - false - } + val useMultipleValuesPerKey = SchemaUtil.checkVariableType(stateVariableInfoOpt, + StateVariableType.ListState) val provider = StateStoreProvider.createAndInit( stateStoreProviderId, keySchema, valueSchema, keyStateEncoderSpec, @@ -149,7 +146,7 @@ abstract class StatePartitionReaderBase( /** * An implementation of [[StatePartitionReaderBase]] for the normal mode of State Data - * Source. It reads the the state at a particular batchId. + * Source. It reads the state at a particular batchId. */ class StatePartitionReader( storeConf: StateStoreConf, @@ -181,41 +178,17 @@ class StatePartitionReader( override lazy val iter: Iterator[InternalRow] = { val stateVarName = stateVariableInfoOpt .map(_.stateName).getOrElse(StateStore.DEFAULT_COL_FAMILY_NAME) - if (SchemaUtil.isMapStateVariable(stateVariableInfoOpt)) { - SchemaUtil.unifyMapStateRowPair( - store.iterator(stateVarName), keySchema, partition.partition) + + if (stateVariableInfoOpt.isDefined) { + val stateVariableInfo = stateVariableInfoOpt.get + val stateVarType = stateVariableInfo.stateVariableType + SchemaUtil.processStateEntries(stateVarType, stateVarName, store, + keySchema, partition.partition, partition.sourceOptions) } else { store .iterator(stateVarName) .map { pair => - stateVariableInfoOpt match { - case Some(stateVarInfo) => - val stateVarType = stateVarInfo.stateVariableType - - stateVarType match { - case StateVariableType.ValueState => - SchemaUtil.unifyStateRowPair((pair.key, pair.value), partition.partition) - - case StateVariableType.ListState => - val key = pair.key - val result = store.valuesIterator(key, stateVarName) - var unsafeRowArr: Seq[UnsafeRow] = Seq.empty - result.foreach { entry => - unsafeRowArr = unsafeRowArr :+ entry.copy() - } - // convert the list of values to array type - val arrData = new GenericArrayData(unsafeRowArr.toArray) - SchemaUtil.unifyStateRowPairWithMultipleValues((pair.key, arrData), - partition.partition) - - case _ => - throw new IllegalStateException( - s"Unsupported state variable type: $stateVarType") - } - - case None => - SchemaUtil.unifyStateRowPair((pair.key, pair.value), partition.partition) - } + SchemaUtil.unifyStateRowPair((pair.key, pair.value), partition.partition) } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala index 88ea06d598e56..dc0d6af951143 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.execution.datasources.v2.state.utils import scala.collection.mutable +import scala.collection.mutable.ArrayBuffer import scala.util.control.NonFatal import org.apache.spark.sql.AnalysisException @@ -24,9 +25,9 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.sql.execution.datasources.v2.state.{StateDataSourceErrors, StateSourceOptions} +import org.apache.spark.sql.execution.streaming.{StateVariableType, TransformWithStateVariableInfo} import org.apache.spark.sql.execution.streaming.StateVariableType._ -import org.apache.spark.sql.execution.streaming.TransformWithStateVariableInfo -import org.apache.spark.sql.execution.streaming.state.{StateStoreColFamilySchema, UnsafeRowPair} +import org.apache.spark.sql.execution.streaming.state.{ReadStateStore, StateStoreColFamilySchema, UnsafeRowPair} import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructType} import org.apache.spark.util.ArrayImplicits._ @@ -58,7 +59,7 @@ object SchemaUtil { } else if (transformWithStateVariableInfoOpt.isDefined) { require(stateStoreColFamilySchemaOpt.isDefined) generateSchemaForStateVar(transformWithStateVariableInfoOpt.get, - stateStoreColFamilySchemaOpt.get) + stateStoreColFamilySchemaOpt.get, sourceOptions) } else { new StructType() .add("key", keySchema) @@ -101,7 +102,8 @@ object SchemaUtil { def unifyMapStateRowPair( stateRows: Iterator[UnsafeRowPair], compositeKeySchema: StructType, - partitionId: Int): Iterator[InternalRow] = { + partitionId: Int, + stateSourceOptions: StateSourceOptions): Iterator[InternalRow] = { val groupingKeySchema = SchemaUtil.getSchemaAsDataType( compositeKeySchema, "key" ).asInstanceOf[StructType] @@ -130,61 +132,84 @@ object SchemaUtil { row } - // All of the rows with the same grouping key were co-located and were - // grouped together consecutively. - new Iterator[InternalRow] { - var curGroupingKey: UnsafeRow = _ - var curStateRowPair: UnsafeRowPair = _ - val curMap = mutable.Map.empty[Any, Any] - - override def hasNext: Boolean = - stateRows.hasNext || !curMap.isEmpty - - override def next(): InternalRow = { - var foundNewGroupingKey = false - while (stateRows.hasNext && !foundNewGroupingKey) { - curStateRowPair = stateRows.next() - if (curGroupingKey == null) { - // First time in the iterator - // Need to make a copy because we need to keep the - // value across function calls - curGroupingKey = curStateRowPair.key - .get(0, groupingKeySchema).asInstanceOf[UnsafeRow].copy() - appendKVPairToMap(curMap, curStateRowPair) - } else { - val curPairGroupingKey = - curStateRowPair.key.get(0, groupingKeySchema) - if (curPairGroupingKey == curGroupingKey) { + def createFlattenedRow( + groupingKey: UnsafeRow, + userMapKey: UnsafeRow, + userMapValue: UnsafeRow, + partitionId: Int): GenericInternalRow = { + val row = new GenericInternalRow(4) + row.update(0, groupingKey) + row.update(1, userMapKey) + row.update(2, userMapValue) + row.update(3, partitionId) + row + } + + if (stateSourceOptions.flattenCollectionTypes) { + stateRows + .map { pair => + val groupingKey = pair.key.get(0, groupingKeySchema).asInstanceOf[UnsafeRow] + val userMapKey = pair.key.get(1, userKeySchema).asInstanceOf[UnsafeRow] + val userMapValue = pair.value + createFlattenedRow(groupingKey, userMapKey, userMapValue, partitionId) + } + } else { + // All of the rows with the same grouping key were co-located and were + // grouped together consecutively. + new Iterator[InternalRow] { + var curGroupingKey: UnsafeRow = _ + var curStateRowPair: UnsafeRowPair = _ + val curMap = mutable.Map.empty[Any, Any] + + override def hasNext: Boolean = + stateRows.hasNext || !curMap.isEmpty + + override def next(): InternalRow = { + var foundNewGroupingKey = false + while (stateRows.hasNext && !foundNewGroupingKey) { + curStateRowPair = stateRows.next() + if (curGroupingKey == null) { + // First time in the iterator + // Need to make a copy because we need to keep the + // value across function calls + curGroupingKey = curStateRowPair.key + .get(0, groupingKeySchema).asInstanceOf[UnsafeRow].copy() appendKVPairToMap(curMap, curStateRowPair) } else { - // find a different grouping key, exit loop and return a row - foundNewGroupingKey = true + val curPairGroupingKey = + curStateRowPair.key.get(0, groupingKeySchema) + if (curPairGroupingKey == curGroupingKey) { + appendKVPairToMap(curMap, curStateRowPair) + } else { + // find a different grouping key, exit loop and return a row + foundNewGroupingKey = true + } } } - } - if (foundNewGroupingKey) { - // found a different grouping key - val row = createDataRow(curGroupingKey, curMap) - // update vars - curGroupingKey = - curStateRowPair.key.get(0, groupingKeySchema) - .asInstanceOf[UnsafeRow].copy() - // empty the map, append current row - curMap.clear() - appendKVPairToMap(curMap, curStateRowPair) - // return map value of previous grouping key - row - } else { - if (curMap.isEmpty) { - throw new NoSuchElementException("Please check if the iterator hasNext(); Likely " + - "user is trying to get element from an exhausted iterator.") - } - else { - // reach the end of the state rows + if (foundNewGroupingKey) { + // found a different grouping key val row = createDataRow(curGroupingKey, curMap) - // clear the map to end the iterator + // update vars + curGroupingKey = + curStateRowPair.key.get(0, groupingKeySchema) + .asInstanceOf[UnsafeRow].copy() + // empty the map, append current row curMap.clear() + appendKVPairToMap(curMap, curStateRowPair) + // return map value of previous grouping key row + } else { + if (curMap.isEmpty) { + throw new NoSuchElementException("Please check if the iterator hasNext(); Likely " + + "user is trying to get element from an exhausted iterator.") + } + else { + // reach the end of the state rows + val row = createDataRow(curGroupingKey, curMap) + // clear the map to end the iterator + curMap.clear() + row + } } } } @@ -200,9 +225,11 @@ object SchemaUtil { "change_type" -> classOf[StringType], "key" -> classOf[StructType], "value" -> classOf[StructType], - "single_value" -> classOf[StructType], + "list_element" -> classOf[StructType], "list_value" -> classOf[ArrayType], "map_value" -> classOf[MapType], + "user_map_key" -> classOf[StructType], + "user_map_value" -> classOf[StructType], "partition_id" -> classOf[IntegerType]) val expectedFieldNames = if (sourceOptions.readChangeFeed) { @@ -213,13 +240,21 @@ object SchemaUtil { stateVarType match { case ValueState => - Seq("key", "single_value", "partition_id") + Seq("key", "value", "partition_id") case ListState => - Seq("key", "list_value", "partition_id") + if (sourceOptions.flattenCollectionTypes) { + Seq("key", "list_element", "partition_id") + } else { + Seq("key", "list_value", "partition_id") + } case MapState => - Seq("key", "map_value", "partition_id") + if (sourceOptions.flattenCollectionTypes) { + Seq("key", "user_map_key", "user_map_value", "partition_id") + } else { + Seq("key", "map_value", "partition_id") + } case _ => throw StateDataSourceErrors @@ -241,21 +276,29 @@ object SchemaUtil { private def generateSchemaForStateVar( stateVarInfo: TransformWithStateVariableInfo, - stateStoreColFamilySchema: StateStoreColFamilySchema): StructType = { + stateStoreColFamilySchema: StateStoreColFamilySchema, + stateSourceOptions: StateSourceOptions): StructType = { val stateVarType = stateVarInfo.stateVariableType stateVarType match { case ValueState => new StructType() .add("key", stateStoreColFamilySchema.keySchema) - .add("single_value", stateStoreColFamilySchema.valueSchema) + .add("value", stateStoreColFamilySchema.valueSchema) .add("partition_id", IntegerType) case ListState => - new StructType() - .add("key", stateStoreColFamilySchema.keySchema) - .add("list_value", ArrayType(stateStoreColFamilySchema.valueSchema)) - .add("partition_id", IntegerType) + if (stateSourceOptions.flattenCollectionTypes) { + new StructType() + .add("key", stateStoreColFamilySchema.keySchema) + .add("list_element", stateStoreColFamilySchema.valueSchema) + .add("partition_id", IntegerType) + } else { + new StructType() + .add("key", stateStoreColFamilySchema.keySchema) + .add("list_value", ArrayType(stateStoreColFamilySchema.valueSchema)) + .add("partition_id", IntegerType) + } case MapState => val groupingKeySchema = SchemaUtil.getSchemaAsDataType( @@ -266,43 +309,47 @@ object SchemaUtil { valueType = stateStoreColFamilySchema.valueSchema ) - new StructType() - .add("key", groupingKeySchema) - .add("map_value", valueMapSchema) - .add("partition_id", IntegerType) + if (stateSourceOptions.flattenCollectionTypes) { + new StructType() + .add("key", groupingKeySchema) + .add("user_map_key", userKeySchema) + .add("user_map_value", stateStoreColFamilySchema.valueSchema) + .add("partition_id", IntegerType) + } else { + new StructType() + .add("key", groupingKeySchema) + .add("map_value", valueMapSchema) + .add("partition_id", IntegerType) + } case _ => throw StateDataSourceErrors.internalError(s"Unsupported state variable type $stateVarType") } } - /** - * Helper functions for map state data source reader. - * - * Map state variables are stored in RocksDB state store has the schema of - * `TransformWithStateKeyValueRowSchemaUtils.getCompositeKeySchema()`; - * But for state store reader, we need to return in format of: - * "key": groupingKey, "map_value": Map(userKey -> value). - * - * The following functions help to translate between two schema. - */ - def isMapStateVariable( - stateVariableInfoOpt: Option[TransformWithStateVariableInfo]): Boolean = { + def checkVariableType( + stateVariableInfoOpt: Option[TransformWithStateVariableInfo], + varType: StateVariableType): Boolean = { stateVariableInfoOpt.isDefined && - stateVariableInfoOpt.get.stateVariableType == MapState + stateVariableInfoOpt.get.stateVariableType == varType } /** * Given key-value schema generated from `generateSchemaForStateVar()`, * returns the compositeKey schema that key is stored in the state store */ - def getCompositeKeySchema(schema: StructType): StructType = { + def getCompositeKeySchema( + schema: StructType, + stateSourceOptions: StateSourceOptions): StructType = { val groupingKeySchema = SchemaUtil.getSchemaAsDataType( schema, "key").asInstanceOf[StructType] val userKeySchema = try { - Option( - SchemaUtil.getSchemaAsDataType(schema, "map_value").asInstanceOf[MapType] + if (stateSourceOptions.flattenCollectionTypes) { + Option(SchemaUtil.getSchemaAsDataType(schema, "user_map_key").asInstanceOf[StructType]) + } else { + Option(SchemaUtil.getSchemaAsDataType(schema, "map_value").asInstanceOf[MapType] .keyType.asInstanceOf[StructType]) + } } catch { case NonFatal(e) => throw StateDataSourceErrors.internalError(s"No such field named as 'map_value' " + @@ -312,4 +359,57 @@ object SchemaUtil { .add("key", groupingKeySchema) .add("userKey", userKeySchema.get) } + + def processStateEntries( + stateVarType: StateVariableType, + stateVarName: String, + store: ReadStateStore, + compositeKeySchema: StructType, + partitionId: Int, + stateSourceOptions: StateSourceOptions): Iterator[InternalRow] = { + stateVarType match { + case StateVariableType.ValueState => + store + .iterator(stateVarName) + .map { pair => + unifyStateRowPair((pair.key, pair.value), partitionId) + } + + case StateVariableType.ListState => + if (stateSourceOptions.flattenCollectionTypes) { + store + .iterator(stateVarName) + .flatMap { pair => + val key = pair.key + val result = store.valuesIterator(key, stateVarName) + result.map { entry => + SchemaUtil.unifyStateRowPair((key, entry), partitionId) + } + } + } else { + store + .iterator(stateVarName) + .map { pair => + val key = pair.key + val result = store.valuesIterator(key, stateVarName) + val unsafeRowArr = ArrayBuffer[UnsafeRow]() + result.foreach { entry => + unsafeRowArr += entry.copy() + } + // convert the list of values to array type + val arrData = new GenericArrayData(unsafeRowArr.toArray) + // convert the list of values to a single row + SchemaUtil.unifyStateRowPairWithMultipleValues((key, arrData), partitionId) + } + } + + case StateVariableType.MapState => + unifyMapStateRowPair(store.iterator(stateVarName), + compositeKeySchema, partitionId, stateSourceOptions) + + case _ => + throw new IllegalStateException( + s"Unsupported state variable type: $stateVarType") + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala index af07707569500..8707facc4c126 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala @@ -287,6 +287,25 @@ class StateDataSourceNegativeTestSuite extends StateDataSourceTestBase { matchPVals = true) } } + + test("ERROR: trying to specify non boolean value for " + + "flattenCollectionTypes") { + withTempDir { tempDir => + runDropDuplicatesQuery(tempDir.getAbsolutePath) + + val exc = intercept[StateDataSourceInvalidOptionValue] { + spark.read.format("statestore") + // trick to bypass getting the last committed batch before validating operator ID + .option(StateSourceOptions.BATCH_ID, 0) + .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, "test") + .load(tempDir.getAbsolutePath) + } + checkError(exc, "STDS_INVALID_OPTION_VALUE.WITH_MESSAGE", Some("42616"), + Map("optionName" -> StateSourceOptions.FLATTEN_COLLECTION_TYPES, + "message" -> ".*"), + matchPVals = true) + } + } } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala index 61091fde35e79..69df86fd5f746 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala @@ -159,7 +159,7 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest val resultDf = stateReaderDf.selectExpr( "key.value AS groupingKey", - "single_value.id AS valueId", "single_value.name AS valueName", + "value.id AS valueId", "value.name AS valueName", "partition_id") checkAnswer(resultDf, @@ -222,7 +222,7 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest .load() val resultDf = stateReaderDf.selectExpr( - "key.value", "single_value.value", "single_value.ttlExpirationMs", "partition_id") + "key.value", "value.value", "value.ttlExpirationMs", "partition_id") var count = 0L resultDf.collect().foreach { row => @@ -235,7 +235,7 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest val answerDf = stateReaderDf.selectExpr( "key.value AS groupingKey", - "single_value.value.value AS valueId", "partition_id") + "value.value.value AS valueId", "partition_id") checkAnswer(answerDf, Seq(Row("a", 1L, 0), Row("b", 1L, 1))) @@ -290,10 +290,12 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest StopStream ) + // Verify that the state can be read in flattened/non-flattened modes val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) .option(StateSourceOptions.STATE_VAR_NAME, "groupsList") + .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, false) .load() val listStateDf = stateReaderDf @@ -307,6 +309,19 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest checkAnswer(listStateDf, Seq(Row("session1", "group1"), Row("session1", "group2"), Row("session1", "group4"), Row("session2", "group1"), Row("session3", "group7"))) + + val flattenedReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "groupsList") + .load() + + val resultDf = flattenedReaderDf.selectExpr( + "key.value AS groupingKey", + "list_element.value AS valueList") + checkAnswer(resultDf, + Seq(Row("session1", "group1"), Row("session1", "group2"), Row("session1", "group4"), + Row("session2", "group1"), Row("session3", "group7"))) } } } @@ -338,10 +353,12 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest StopStream ) + // Verify that the state can be read in flattened/non-flattened modes val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) .option(StateSourceOptions.STATE_VAR_NAME, "groupsListWithTTL") + .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, false) .load() val listStateDf = stateReaderDf @@ -368,6 +385,31 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest checkAnswer(valuesDf, Seq(Row("session1", "group1"), Row("session1", "group2"), Row("session1", "group4"), Row("session2", "group1"), Row("session3", "group7"))) + + val flattenedStateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "groupsListWithTTL") + .load() + + val flattenedResultDf = flattenedStateReaderDf + .selectExpr("list_element.ttlExpirationMs AS ttlExpirationMs") + var flattenedCount = 0L + flattenedResultDf.collect().foreach { row => + flattenedCount = flattenedCount + 1 + assert(row.getLong(0) > 0) + } + + // verify that 5 state rows are present + assert(flattenedCount === 5) + + val outputDf = flattenedStateReaderDf + .selectExpr("key.value AS groupingKey", + "list_element.value.value AS groupId") + + checkAnswer(outputDf, + Seq(Row("session1", "group1"), Row("session1", "group2"), Row("session1", "group4"), + Row("session2", "group1"), Row("session3", "group7"))) } } } @@ -397,10 +439,12 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest StopStream ) + // Verify that the state can be read in flattened/non-flattened modes val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) .option(StateSourceOptions.STATE_VAR_NAME, "sessionState") + .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, false) .load() val resultDf = stateReaderDf.selectExpr( @@ -413,6 +457,24 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest Row("k2", Map(Row("v2") -> Row("3")))) ) + + val flattenedStateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "sessionState") + .load() + + val outputDf = flattenedStateReaderDf + .selectExpr("key.value AS groupingKey", + "user_map_key.value AS mapKey", + "user_map_value.value AS mapValue") + + checkAnswer(outputDf, + Seq( + Row("k1", "v1", "10"), + Row("k1", "v2", "5"), + Row("k2", "v2", "3")) + ) } } } @@ -463,10 +525,12 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest StopStream ) + // Verify that the state can be read in flattened/non-flattened modes val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) .option(StateSourceOptions.STATE_VAR_NAME, "mapState") + .option(StateSourceOptions.FLATTEN_COLLECTION_TYPES, false) .load() val resultDf = stateReaderDf.selectExpr( @@ -478,6 +542,24 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest Map(Row("key2") -> Row(Row(2), 61000L), Row("key1") -> Row(Row(1), 61000L)))) ) + + val flattenedStateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "mapState") + .load() + + val outputDf = flattenedStateReaderDf + .selectExpr("key.value AS groupingKey", + "user_map_key.value AS mapKey", + "user_map_value.value.value AS mapValue", + "user_map_value.ttlExpirationMs AS ttlTimestamp") + + checkAnswer(outputDf, + Seq( + Row("k1", "key1", 1, 61000L), + Row("k1", "key2", 2, 61000L)) + ) } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala index d0e255bb30499..0c02fbf97820b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala @@ -1623,7 +1623,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest val batch1AnsDf = batch1Df.selectExpr( "key.value AS groupingKey", - "single_value.value AS valueId") + "value.value AS valueId") checkAnswer(batch1AnsDf, Seq(Row("a", 2L))) @@ -1636,7 +1636,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest val batch3AnsDf = batch3Df.selectExpr( "key.value AS groupingKey", - "single_value.value AS valueId") + "value.value AS valueId") checkAnswer(batch3AnsDf, Seq(Row("a", 1L))) } } @@ -1731,7 +1731,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest val countStateAnsDf = countStateDf.selectExpr( "key.value AS groupingKey", - "single_value.value AS valueId") + "value.value AS valueId") checkAnswer(countStateAnsDf, Seq(Row("a", 5L))) val mostRecentDf = spark.read @@ -1743,7 +1743,7 @@ class TransformWithStateSuite extends StateStoreMetricsTest val mostRecentAnsDf = mostRecentDf.selectExpr( "key.value AS groupingKey", - "single_value.value") + "value.value") checkAnswer(mostRecentAnsDf, Seq(Row("a", "str1"))) } } From dedf5aa91827f32736ce5dae2eb123ba4e244c3b Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Tue, 24 Sep 2024 07:40:58 -0700 Subject: [PATCH 119/189] [SPARK-49750][DOC] Mention delegation token support in K8s mode ### What changes were proposed in this pull request? Update docs to mention delegation token support in K8s mode. ### Why are the changes needed? The delegation token support in K8s mode has been implemented since 3.0.0 via SPARK-23257. ### Does this PR introduce _any_ user-facing change? Yes, docs are updated. ### How was this patch tested? Review. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48199 from pan3793/SPARK-49750. Authored-by: Cheng Pan Signed-off-by: Dongjoon Hyun --- docs/security.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/security.md b/docs/security.md index b97abfeacf240..c7d3fd5f8c36f 100644 --- a/docs/security.md +++ b/docs/security.md @@ -947,7 +947,7 @@ mechanism (see `java.util.ServiceLoader`). Implementations of `org.apache.spark.security.HadoopDelegationTokenProvider` can be made available to Spark by listing their names in the corresponding file in the jar's `META-INF/services` directory. -Delegation token support is currently only supported in YARN mode. Consult the +Delegation token support is currently only supported in YARN and Kubernetes mode. Consult the deployment-specific page for more information. The following options provides finer-grained control for this feature: From 55d0233d19cc52bee91a9619057d9b6f33165a0a Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 24 Sep 2024 07:48:23 -0700 Subject: [PATCH 120/189] [SPARK-49713][PYTHON][FOLLOWUP] Make function `count_min_sketch` accept long seed ### What changes were proposed in this pull request? Make function `count_min_sketch` accept long seed ### Why are the changes needed? existing implementation only accepts int seed, which is inconsistent with other `ExpressionWithRandomSeed`: ```py In [3]: >>> from pyspark.sql import functions as sf ...: >>> spark.range(100).select( ...: ... sf.hex(sf.count_min_sketch("id", sf.lit(1.5), 0.6, 1111111111111111111)) ...: ... ).show(truncate=False) ... AnalysisException: [DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "count_min_sketch(id, 1.5, 0.6, 1111111111111111111)" due to data type mismatch: The 4th parameter requires the "INT" type, however "1111111111111111111" has the type "BIGINT". SQLSTATE: 42K09; 'Aggregate [unresolvedalias('hex(count_min_sketch(id#64L, 1.5, 0.6, 1111111111111111111, 0, 0)))] +- Range (0, 100, step=1, splits=Some(12)) ... ``` ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? added doctest ### Was this patch authored or co-authored using generative AI tooling? no Closes #48223 from zhengruifeng/count_min_sk_long_seed. Authored-by: Ruifeng Zheng Signed-off-by: Dongjoon Hyun --- python/pyspark/sql/connect/functions/builtin.py | 3 +-- python/pyspark/sql/functions/builtin.py | 14 +++++++++++++- .../scala/org/apache/spark/sql/functions.scala | 2 +- .../expressions/aggregate/CountMinSketchAgg.scala | 8 ++++++-- 4 files changed, 21 insertions(+), 6 deletions(-) diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index 2a39bc6bfddda..6953230f5b42e 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -70,7 +70,6 @@ StringType, ) from pyspark.sql.utils import enum_to_value as _enum_to_value -from pyspark.util import JVM_INT_MAX # The implementation of pandas_udf is embedded in pyspark.sql.function.pandas_udf # for code reuse. @@ -1130,7 +1129,7 @@ def count_min_sketch( confidence: Union[Column, float], seed: Optional[Union[Column, int]] = None, ) -> Column: - _seed = lit(random.randint(0, JVM_INT_MAX)) if seed is None else lit(seed) + _seed = lit(random.randint(0, sys.maxsize)) if seed is None else lit(seed) return _invoke_function_over_columns("count_min_sketch", col, lit(eps), lit(confidence), _seed) diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 2688f9daa23a4..09a286fe7c94e 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -6080,7 +6080,19 @@ def count_min_sketch( |0000000100000000000000640000000100000002000000005D96391C00000000000000320000000000000032| +----------------------------------------------------------------------------------------+ - Example 3: Using a random seed + Example 3: Using a long seed + + >>> from pyspark.sql import functions as sf + >>> spark.range(100).select( + ... sf.hex(sf.count_min_sketch("id", sf.lit(1.5), 0.2, 1111111111111111111)) + ... ).show(truncate=False) + +----------------------------------------------------------------------------------------+ + |hex(count_min_sketch(id, 1.5, 0.2, 1111111111111111111)) | + +----------------------------------------------------------------------------------------+ + |00000001000000000000006400000001000000020000000044078BA100000000000000320000000000000032| + +----------------------------------------------------------------------------------------+ + + Example 4: Using a random seed >>> from pyspark.sql import functions as sf >>> spark.range(100).select( diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index d9bceabe88f8f..ab69789c75f50 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -399,7 +399,7 @@ object functions { * @since 4.0.0 */ def count_min_sketch(e: Column, eps: Column, confidence: Column): Column = - count_min_sketch(e, eps, confidence, lit(SparkClassUtils.random.nextInt)) + count_min_sketch(e, eps, confidence, lit(SparkClassUtils.random.nextLong)) private[spark] def collect_top_k(e: Column, num: Int, reverse: Boolean): Column = Column.internalFn("collect_top_k", e, lit(num), lit(reverse)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala index c26c4a9bdfea3..f0a27677628dc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala @@ -63,7 +63,10 @@ case class CountMinSketchAgg( // Mark as lazy so that they are not evaluated during tree transformation. private lazy val eps: Double = epsExpression.eval().asInstanceOf[Double] private lazy val confidence: Double = confidenceExpression.eval().asInstanceOf[Double] - private lazy val seed: Int = seedExpression.eval().asInstanceOf[Int] + private lazy val seed: Int = seedExpression.eval() match { + case i: Int => i + case l: Long => l.toInt + } override def checkInputDataTypes(): TypeCheckResult = { val defaultCheck = super.checkInputDataTypes() @@ -168,7 +171,8 @@ case class CountMinSketchAgg( copy(inputAggBufferOffset = newInputAggBufferOffset) override def inputTypes: Seq[AbstractDataType] = { - Seq(TypeCollection(IntegralType, StringType, BinaryType), DoubleType, DoubleType, IntegerType) + Seq(TypeCollection(IntegralType, StringType, BinaryType), DoubleType, DoubleType, + TypeCollection(IntegerType, LongType)) } override def nullable: Boolean = false From afe8bf945e1ad72fcb0ec4ec35b169e54169f5f1 Mon Sep 17 00:00:00 2001 From: allisonwang-db Date: Wed, 25 Sep 2024 08:53:09 +0900 Subject: [PATCH 121/189] [SPARK-49771][PYTHON] Improve Pandas Scalar Iter UDF error when output rows exceed input rows ### What changes were proposed in this pull request? This PR changes the `assert` error into a user-facing PySpark error when the pandas_iter UDF has more output rows than input rows. ### Why are the changes needed? To make the error message more user-friendly. After the PR, the error will be `pyspark.errors.exceptions.base.PySparkRuntimeError: [PANDAS_UDF_OUTPUT_EXCEEDS_INPUT_ROWS] The Pandas SCALAR_ITER UDF outputs more rows than input rows.` ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #48231 from allisonwang-db/spark-49771-pd-iter-err. Authored-by: allisonwang-db Signed-off-by: Hyukjin Kwon --- python/pyspark/errors/error-conditions.json | 5 +++++ python/pyspark/worker.py | 9 +++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 92aeb15e21d1b..115ad658e32f5 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -802,6 +802,11 @@ " >= must be installed; however, it was not found." ] }, + "PANDAS_UDF_OUTPUT_EXCEEDS_INPUT_ROWS" : { + "message": [ + "The Pandas SCALAR_ITER UDF outputs more rows than input rows." + ] + }, "PIPE_FUNCTION_EXITED": { "message": [ "Pipe function `` exited with error code ." diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index b8263769c28a9..eedf5d1fd5996 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -1565,14 +1565,15 @@ def map_batch(batch): num_output_rows = 0 for result_batch, result_type in result_iter: num_output_rows += len(result_batch) - # This assert is for Scalar Iterator UDF to fail fast. + # This check is for Scalar Iterator UDF to fail fast. # The length of the entire input can only be explicitly known # by consuming the input iterator in user side. Therefore, # it's very unlikely the output length is higher than # input length. - assert ( - is_map_pandas_iter or is_map_arrow_iter or num_output_rows <= num_input_rows - ), "Pandas SCALAR_ITER UDF outputted more rows than input rows." + if is_scalar_iter and num_output_rows > num_input_rows: + raise PySparkRuntimeError( + errorClass="PANDAS_UDF_OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={} + ) yield (result_batch, result_type) if is_scalar_iter: From 0a7b98532fd2cf3a251aa258886c1e78779e9594 Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Wed, 25 Sep 2024 08:57:55 +0900 Subject: [PATCH 122/189] [SPARK-49585][CONNECT] Replace executions map in SessionHolder with operationID set ### What changes were proposed in this pull request? SessionHolder has no reason to store ExecuteHolder directly as SparkConnectExecutionManager has a global map of ExecuteHolder. This PR replaces the map in SessionHolder with a set of operation IDs which is only used when interrupting executions within the session. ### Why are the changes needed? Save memory, and simplify the code by making SparkConnectExecutionManager the single source of ExecuteHolders. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48071 from changgyoopark-db/SPARK-49585. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../sql/connect/service/SessionHolder.scala | 62 +++++++++---------- .../SparkConnectExecutionManager.scala | 4 +- .../SparkConnectReattachExecuteHandler.scala | 33 +++++----- .../SparkConnectReleaseExecuteHandler.scala | 4 +- .../planner/SparkConnectServiceSuite.scala | 6 +- .../SparkConnectSessionHolderSuite.scala | 9 +++ 6 files changed, 67 insertions(+), 51 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala index e56d66da3050d..5dced7acfb0d2 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala @@ -19,7 +19,6 @@ package org.apache.spark.sql.connect.service import java.nio.file.Path import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, TimeUnit} -import javax.annotation.concurrent.GuardedBy import scala.collection.mutable import scala.concurrent.{ExecutionContext, Future} @@ -40,6 +39,7 @@ import org.apache.spark.sql.connect.common.InvalidPlanInput import org.apache.spark.sql.connect.config.Connect import org.apache.spark.sql.connect.planner.PythonStreamingQueryListener import org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper +import org.apache.spark.sql.connect.service.ExecuteKey import org.apache.spark.sql.connect.service.SessionHolder.{ERROR_CACHE_SIZE, ERROR_CACHE_TIMEOUT_SEC} import org.apache.spark.sql.streaming.StreamingQueryListener import org.apache.spark.util.{SystemClock, Utils} @@ -91,8 +91,8 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio // Setting it to -1 indicated forever. @volatile private var customInactiveTimeoutMs: Option[Long] = None - private val executions: ConcurrentMap[String, ExecuteHolder] = - new ConcurrentHashMap[String, ExecuteHolder]() + private val operationIds: ConcurrentMap[String, Boolean] = + new ConcurrentHashMap[String, Boolean]() // The cache that maps an error id to a throwable. The throwable in cache is independent to // each other. @@ -138,12 +138,11 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio } /** - * Add ExecuteHolder to this session. + * Add an operation ID to this session. * - * Called only by SparkConnectExecutionManager under executionsLock. + * Called only by SparkConnectExecutionManager when a new execution is started. */ - @GuardedBy("SparkConnectService.executionManager.executionsLock") - private[service] def addExecuteHolder(executeHolder: ExecuteHolder): Unit = { + private[service] def addOperationId(operationId: String): Unit = { if (closedTimeMs.isDefined) { // Do not accept new executions if the session is closing. throw new SparkSQLException( @@ -151,26 +150,20 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio messageParameters = Map("handle" -> sessionId)) } - val oldExecute = executions.putIfAbsent(executeHolder.operationId, executeHolder) - if (oldExecute != null) { - // the existence of this should alrady be checked by SparkConnectExecutionManager - throw new IllegalStateException( - s"ExecuteHolder with opId=${executeHolder.operationId} already exists!") + val alreadyExists = operationIds.putIfAbsent(operationId, true) + if (alreadyExists) { + // The existence of it should have been checked by SparkConnectExecutionManager. + throw new IllegalStateException(s"ExecuteHolder with opId=${operationId} already exists!") } } /** - * Remove ExecuteHolder from this session. + * Remove an operation ID from this session. * - * Called only by SparkConnectExecutionManager under executionsLock. + * Called only by SparkConnectExecutionManager when an execution is ended. */ - @GuardedBy("SparkConnectService.executionManager.executionsLock") - private[service] def removeExecuteHolder(operationId: String): Unit = { - executions.remove(operationId) - } - - private[connect] def executeHolder(operationId: String): Option[ExecuteHolder] = { - Option(executions.get(operationId)) + private[service] def removeOperationId(operationId: String): Unit = { + operationIds.remove(operationId) } /** @@ -182,9 +175,12 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio val interruptedIds = new mutable.ArrayBuffer[String]() val operationsIds = SparkConnectService.streamingSessionManager.cleanupRunningQueries(this, blocking = false) - executions.asScala.values.foreach { execute => - if (execute.interrupt()) { - interruptedIds += execute.operationId + operationIds.asScala.foreach { case (operationId, _) => + val executeKey = ExecuteKey(userId, sessionId, operationId) + SparkConnectService.executionManager.getExecuteHolder(executeKey).foreach { executeHolder => + if (executeHolder.interrupt()) { + interruptedIds += operationId + } } } interruptedIds.toSeq ++ operationsIds @@ -199,10 +195,13 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio val interruptedIds = new mutable.ArrayBuffer[String]() val queries = SparkConnectService.streamingSessionManager.getTaggedQuery(tag, session) queries.foreach(q => Future(q.query.stop())(ExecutionContext.global)) - executions.asScala.values.foreach { execute => - if (execute.sparkSessionTags.contains(tag)) { - if (execute.interrupt()) { - interruptedIds += execute.operationId + operationIds.asScala.foreach { case (operationId, _) => + val executeKey = ExecuteKey(userId, sessionId, operationId) + SparkConnectService.executionManager.getExecuteHolder(executeKey).foreach { executeHolder => + if (executeHolder.sparkSessionTags.contains(tag)) { + if (executeHolder.interrupt()) { + interruptedIds += operationId + } } } } @@ -216,9 +215,10 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio */ private[service] def interruptOperation(operationId: String): Seq[String] = { val interruptedIds = new mutable.ArrayBuffer[String]() - Option(executions.get(operationId)).foreach { execute => - if (execute.interrupt()) { - interruptedIds += execute.operationId + val executeKey = ExecuteKey(userId, sessionId, operationId) + SparkConnectService.executionManager.getExecuteHolder(executeKey).foreach { executeHolder => + if (executeHolder.interrupt()) { + interruptedIds += operationId } } interruptedIds.toSeq diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala index d66964b8d34bd..d9eb5438c3886 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutionManager.scala @@ -114,7 +114,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { new ExecuteHolder(executeKey, request, sessionHolder) }) - sessionHolder.addExecuteHolder(executeHolder) + sessionHolder.addOperationId(executeHolder.operationId) logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, executeHolder.key)} is created.") @@ -142,7 +142,7 @@ private[connect] class SparkConnectExecutionManager() extends Logging { // Remove the execution from the map *after* putting it in abandonedTombstones. executions.remove(key) - executeHolder.sessionHolder.removeExecuteHolder(executeHolder.operationId) + executeHolder.sessionHolder.removeOperationId(executeHolder.operationId) updateLastExecutionTime() diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReattachExecuteHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReattachExecuteHandler.scala index 534937f84eaee..a2696311bd843 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReattachExecuteHandler.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReattachExecuteHandler.scala @@ -23,6 +23,7 @@ import org.apache.spark.SparkSQLException import org.apache.spark.connect.proto import org.apache.spark.internal.Logging import org.apache.spark.sql.connect.execution.ExecuteGrpcResponseSender +import org.apache.spark.sql.connect.service.ExecuteKey class SparkConnectReattachExecuteHandler( responseObserver: StreamObserver[proto.ExecutePlanResponse]) @@ -38,22 +39,24 @@ class SparkConnectReattachExecuteHandler( SessionKey(v.getUserContext.getUserId, v.getSessionId), previousSessionId) - val executeHolder = sessionHolder.executeHolder(v.getOperationId).getOrElse { - if (SparkConnectService.executionManager - .getAbandonedTombstone( - ExecuteKey(v.getUserContext.getUserId, v.getSessionId, v.getOperationId)) - .isDefined) { - logDebug(s"Reattach operation abandoned: ${v.getOperationId}") - throw new SparkSQLException( - errorClass = "INVALID_HANDLE.OPERATION_ABANDONED", - messageParameters = Map("handle" -> v.getOperationId)) - } else { - logDebug(s"Reattach operation not found: ${v.getOperationId}") - throw new SparkSQLException( - errorClass = "INVALID_HANDLE.OPERATION_NOT_FOUND", - messageParameters = Map("handle" -> v.getOperationId)) + val executeKey = ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, v.getOperationId) + val executeHolder = + SparkConnectService.executionManager.getExecuteHolder(executeKey).getOrElse { + if (SparkConnectService.executionManager + .getAbandonedTombstone( + ExecuteKey(v.getUserContext.getUserId, v.getSessionId, v.getOperationId)) + .isDefined) { + logDebug(s"Reattach operation abandoned: ${v.getOperationId}") + throw new SparkSQLException( + errorClass = "INVALID_HANDLE.OPERATION_ABANDONED", + messageParameters = Map("handle" -> v.getOperationId)) + } else { + logDebug(s"Reattach operation not found: ${v.getOperationId}") + throw new SparkSQLException( + errorClass = "INVALID_HANDLE.OPERATION_NOT_FOUND", + messageParameters = Map("handle" -> v.getOperationId)) + } } - } if (!executeHolder.reattachable) { logWarning(s"Reattach to not reattachable operation.") throw new SparkSQLException( diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReleaseExecuteHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReleaseExecuteHandler.scala index a2dbf3b2eec9f..6beba13d55156 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReleaseExecuteHandler.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectReleaseExecuteHandler.scala @@ -22,6 +22,7 @@ import io.grpc.stub.StreamObserver import org.apache.spark.SparkSQLException import org.apache.spark.connect.proto import org.apache.spark.internal.Logging +import org.apache.spark.sql.connect.service.ExecuteKey class SparkConnectReleaseExecuteHandler( responseObserver: StreamObserver[proto.ReleaseExecuteResponse]) @@ -42,8 +43,9 @@ class SparkConnectReleaseExecuteHandler( // ReleaseExecute arrived after it was abandoned and timed out. // An asynchronous ReleastUntil operation may also arrive after ReleaseAll. // Because of that, make it noop and not fail if the ExecuteHolder is no longer there. + val executeKey = ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, v.getOperationId) val executeHolderOption = - sessionHolder.executeHolder(v.getOperationId).foreach { executeHolder => + SparkConnectService.executionManager.getExecuteHolder(executeKey).foreach { executeHolder => if (!executeHolder.reattachable) { throw new SparkSQLException( errorClass = "INVALID_CURSOR.NOT_REATTACHABLE", diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala index 62146f19328a8..d6d137e6d91aa 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala @@ -48,7 +48,7 @@ import org.apache.spark.sql.connect.dsl.MockRemoteSession import org.apache.spark.sql.connect.dsl.expressions._ import org.apache.spark.sql.connect.dsl.plans._ import org.apache.spark.sql.connect.plugin.SparkConnectPluginRegistry -import org.apache.spark.sql.connect.service.{ExecuteHolder, ExecuteStatus, SessionStatus, SparkConnectAnalyzeHandler, SparkConnectService, SparkListenerConnectOperationStarted} +import org.apache.spark.sql.connect.service.{ExecuteHolder, ExecuteKey, ExecuteStatus, SessionStatus, SparkConnectAnalyzeHandler, SparkConnectService, SparkListenerConnectOperationStarted} import org.apache.spark.sql.connector.catalog.InMemoryPartitionTableCatalog import org.apache.spark.sql.streaming.StreamingQuery import org.apache.spark.sql.test.SharedSparkSession @@ -926,7 +926,9 @@ class SparkConnectServiceSuite semaphoreStarted.release() val sessionHolder = SparkConnectService.getOrCreateIsolatedSession(e.userId, e.sessionId, None) - executeHolder = sessionHolder.executeHolder(e.operationId) + val executeKey = + ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, e.operationId) + executeHolder = SparkConnectService.executionManager.getExecuteHolder(executeKey) case _ => } } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala index ed2f60afb0096..21f84291a2f07 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala @@ -413,4 +413,13 @@ class SparkConnectSessionHolderSuite extends SharedSparkSession { planner.transformRelation(query, cachePlan = true) assertPlanCache(sessionHolder, Some(Set())) } + + test("Test duplicate operation IDs") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + sessionHolder.addOperationId("DUMMY") + val ex = intercept[IllegalStateException] { + sessionHolder.addOperationId("DUMMY") + } + assert(ex.getMessage.contains("already exists")) + } } From 29ed2729492a7af3445b436cf589883e56dd9aee Mon Sep 17 00:00:00 2001 From: Changgyoo Park Date: Wed, 25 Sep 2024 08:58:33 +0900 Subject: [PATCH 123/189] [SPARK-49688][CONNECT] Fix a data race between interrupt and execute plan ### What changes were proposed in this pull request? Get rid of the complicated "promise"-based completion callback mechanism, and introduce a lock-free state machine. The gist is, - A thread can only be interrupted when it is in a certain state: started. - A successful interruption means the interrupted thread must call the completion callback. - Interruption after completion or before starting is prohibited without relying on a mutex. ### Why are the changes needed? Execution can be interrupted before started, thus causing the "closed" message to be omitted. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? SparkConnectServiceSuite. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48208 from changgyoopark-db/SPARK-49688. Authored-by: Changgyoo Park Signed-off-by: Hyukjin Kwon --- .../execution/ExecuteThreadRunner.scala | 224 ++++++++++-------- .../sql/connect/service/ExecuteHolder.scala | 35 +-- .../service/SparkConnectServiceE2ESuite.scala | 2 +- 3 files changed, 151 insertions(+), 110 deletions(-) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala index e75654e2c384f..61be2bc4eb994 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala @@ -17,9 +17,9 @@ package org.apache.spark.sql.connect.execution -import scala.concurrent.{ExecutionContext, Promise} +import java.util.concurrent.atomic.AtomicInteger + import scala.jdk.CollectionConverters._ -import scala.util.Try import scala.util.control.NonFatal import com.google.protobuf.Message @@ -32,7 +32,7 @@ import org.apache.spark.sql.connect.common.ProtoUtils import org.apache.spark.sql.connect.planner.SparkConnectPlanner import org.apache.spark.sql.connect.service.{ExecuteHolder, ExecuteSessionTag, SparkConnectService} import org.apache.spark.sql.connect.utils.ErrorUtils -import org.apache.spark.util.{ThreadUtils, Utils} +import org.apache.spark.util.Utils /** * This class launches the actual execution in an execution thread. The execution pushes the @@ -40,68 +40,70 @@ import org.apache.spark.util.{ThreadUtils, Utils} */ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends Logging { - private val promise: Promise[Unit] = Promise[Unit]() + /** The thread state. */ + private val state: AtomicInteger = new AtomicInteger(ThreadState.notStarted) // The newly created thread will inherit all InheritableThreadLocals used by Spark, // e.g. SparkContext.localProperties. If considering implementing a thread-pool, // forwarding of thread locals needs to be taken into account. - private val executionThread: ExecutionThread = new ExecutionThread(promise) - - private var started: Boolean = false - - private var interrupted: Boolean = false - - private var completed: Boolean = false - - private val lock = new Object - - /** Launches the execution in a background thread, returns immediately. */ - private[connect] def start(): Unit = { - lock.synchronized { - assert(!started) - // Do not start if already interrupted. - if (!interrupted) { - executionThread.start() - started = true - } - } - } + private val executionThread: ExecutionThread = new ExecutionThread() /** - * Register a callback that gets executed after completion/interruption of the execution thread. + * Launches the execution in a background thread, returns immediately. This method is expected + * to be invoked only once for an ExecuteHolder. */ - private[connect] def processOnCompletion(callback: Try[Unit] => Unit): Unit = { - promise.future.onComplete(callback)(ExecuteThreadRunner.namedExecutionContext) + private[connect] def start(): Unit = { + val currentState = state.getAcquire() + if (currentState == ThreadState.notStarted) { + executionThread.start() + } else { + // This assertion does not hold if it is called more than once. + assert(currentState == ThreadState.interrupted) + } } /** - * Interrupt the executing thread. + * Interrupts the execution thread if the execution has been interrupted by this method call. + * * @return - * true if it was not interrupted before, false if it was already interrupted or completed. + * true if the thread is running and interrupted. */ private[connect] def interrupt(): Boolean = { - lock.synchronized { - if (!started && !interrupted) { - // execution thread hasn't started yet, and will not be started. - // handle the interrupted error here directly. - interrupted = true - ErrorUtils.handleError( - "execute", - executeHolder.responseObserver, - executeHolder.sessionHolder.userId, - executeHolder.sessionHolder.sessionId, - Some(executeHolder.eventsManager), - interrupted)(new SparkSQLException("OPERATION_CANCELED", Map.empty)) - true - } else if (!interrupted && !completed) { - // checking completed prevents sending interrupt onError after onCompleted - interrupted = true - executionThread.interrupt() - true + var currentState = state.getAcquire() + while (currentState == ThreadState.notStarted || currentState == ThreadState.started) { + val newState = if (currentState == ThreadState.notStarted) { + ThreadState.interrupted } else { - false + ThreadState.startedInterrupted + } + + val prevState = state.compareAndExchangeRelease(currentState, newState) + if (prevState == currentState) { + if (prevState == ThreadState.notStarted) { + // The execution thread has not been started, or will immediately return because the state + // transition happens at the beginning of executeInternal. + try { + ErrorUtils.handleError( + "execute", + executeHolder.responseObserver, + executeHolder.sessionHolder.userId, + executeHolder.sessionHolder.sessionId, + Some(executeHolder.eventsManager), + true)(new SparkSQLException("OPERATION_CANCELED", Map.empty)) + } finally { + executeHolder.cleanup() + } + } else { + // Interrupt execution. + executionThread.interrupt() + } + return true } + currentState = prevState } + + // Already interrupted, completed, or not started. + false } private def execute(): Unit = { @@ -118,15 +120,8 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends executeHolder.sessionHolder.session.sparkContext.cancelJobsWithTag( executeHolder.jobTag, s"A job with the same tag ${executeHolder.jobTag} has failed.") - // Rely on an internal interrupted flag, because Thread.interrupted() could be cleared, - // and different exceptions like InterruptedException, ClosedByInterruptException etc. - // could be thrown. - if (interrupted) { - throw new SparkSQLException("OPERATION_CANCELED", Map.empty) - } else { - // Rethrown the original error. - throw e - } + // Rethrow the original error. + throw e } finally { executeHolder.sessionHolder.session.sparkContext.removeJobTag(executeHolder.jobTag) SparkConnectService.executionListener.foreach(_.removeJobTag(executeHolder.jobTag)) @@ -139,23 +134,50 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends } } } catch { - ErrorUtils.handleError( - "execute", - executeHolder.responseObserver, - executeHolder.sessionHolder.userId, - executeHolder.sessionHolder.sessionId, - Some(executeHolder.eventsManager), - interrupted) + case e: Throwable if state.getAcquire() != ThreadState.startedInterrupted => + ErrorUtils.handleError( + "execute", + executeHolder.responseObserver, + executeHolder.sessionHolder.userId, + executeHolder.sessionHolder.sessionId, + Some(executeHolder.eventsManager), + false)(e) + } finally { + // Make sure to transition to completed in order to prevent the thread from being interrupted + // afterwards. + var currentState = state.getAcquire() + while (currentState == ThreadState.started || + currentState == ThreadState.startedInterrupted) { + val interrupted = currentState == ThreadState.startedInterrupted + val prevState = state.compareAndExchangeRelease(currentState, ThreadState.completed) + if (prevState == currentState) { + if (interrupted) { + try { + ErrorUtils.handleError( + "execute", + executeHolder.responseObserver, + executeHolder.sessionHolder.userId, + executeHolder.sessionHolder.sessionId, + Some(executeHolder.eventsManager), + true)(new SparkSQLException("OPERATION_CANCELED", Map.empty)) + } finally { + executeHolder.cleanup() + } + } + return + } + currentState = prevState + } } } // Inner executeInternal is wrapped by execute() for error handling. - private def executeInternal() = { - // synchronized - check if already got interrupted while starting. - lock.synchronized { - if (interrupted) { - throw new InterruptedException() - } + private def executeInternal(): Unit = { + val prevState = state.compareAndExchangeRelease(ThreadState.notStarted, ThreadState.started) + if (prevState != ThreadState.notStarted) { + // Silently return, expecting that the caller would handle the interruption. + assert(prevState == ThreadState.interrupted) + return } // `withSession` ensures that session-specific artifacts (such as JARs and class files) are @@ -226,17 +248,14 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends observedMetrics ++ accumulatedInPython)) } - lock.synchronized { - // Synchronized before sending ResultComplete, and up until completing the result stream - // to prevent a situation in which a client of reattachable execution receives - // ResultComplete, and proceeds to send ReleaseExecute, and that triggers an interrupt - // before it finishes. - - if (interrupted) { - // check if it got interrupted at the very last moment - throw new InterruptedException() - } - completed = true // no longer interruptible + // State transition should be atomic to prevent a situation in which a client of reattachable + // execution receives ResultComplete, and proceeds to send ReleaseExecute, and that triggers + // an interrupt before it finishes. Failing to transition to completed means that the thread + // was interrupted, and that will be checked at the end of the execution. + if (state.compareAndExchangeRelease( + ThreadState.started, + ThreadState.completed) == ThreadState.started) { + // Now, the execution cannot be interrupted. // If the request starts a long running iterator (e.g. StreamingQueryListener needs // a long-running iterator to continuously stream back events, it runs in a separate @@ -311,21 +330,36 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends .build() } - private class ExecutionThread(onCompletionPromise: Promise[Unit]) + private class ExecutionThread() extends Thread(s"SparkConnectExecuteThread_opId=${executeHolder.operationId}") { - override def run(): Unit = { - try { - execute() - onCompletionPromise.success(()) - } catch { - case NonFatal(e) => - onCompletionPromise.failure(e) - } - } + override def run(): Unit = execute() } } -private[connect] object ExecuteThreadRunner { - private implicit val namedExecutionContext: ExecutionContext = ExecutionContext - .fromExecutor(ThreadUtils.newDaemonSingleThreadExecutor("SparkConnectExecuteThreadCallback")) +/** + * Defines possible execution thread states. + * + * The state transitions as follows. + * - notStarted -> interrupted. + * - notStarted -> started -> startedInterrupted -> completed. + * - notStarted -> started -> completed. + * + * The thread can only be interrupted if the thread is in the startedInterrupted state. + */ +private object ThreadState { + + /** The thread has not started: transition to interrupted or started. */ + val notStarted: Int = 0 + + /** Execution was interrupted: terminal state. */ + val interrupted: Int = 1 + + /** The thread has started: transition to startedInterrupted or completed. */ + val started: Int = 2 + + /** The thread has started and execution was interrupted: transition to completed. */ + val startedInterrupted: Int = 3 + + /** Execution was completed: terminal state. */ + val completed: Int = 4 } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala index dc349c3e33251..821ddb2c85d58 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.connect.service +import java.util.concurrent.atomic.AtomicBoolean + import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -104,8 +106,8 @@ private[connect] class ExecuteHolder( : mutable.ArrayBuffer[ExecuteGrpcResponseSender[proto.ExecutePlanResponse]] = new mutable.ArrayBuffer[ExecuteGrpcResponseSender[proto.ExecutePlanResponse]]() - /** For testing. Whether the async completion callback is called. */ - @volatile private[connect] var completionCallbackCalled: Boolean = false + /** Indicates whether the cleanup method was called. */ + private[connect] val completionCallbackCalled: AtomicBoolean = new AtomicBoolean(false) /** * Start the execution. The execution is started in a background thread in ExecuteThreadRunner. @@ -227,16 +229,7 @@ private[connect] class ExecuteHolder( def close(): Unit = synchronized { if (closedTimeMs.isEmpty) { // interrupt execution, if still running. - runner.interrupt() - // Do not wait for the execution to finish, clean up resources immediately. - runner.processOnCompletion { _ => - completionCallbackCalled = true - // The execution may not immediately get interrupted, clean up any remaining resources when - // it does. - responseObserver.removeAll() - // post closed to UI - eventsManager.postClosed() - } + val interrupted = runner.interrupt() // interrupt any attached grpcResponseSenders grpcResponseSenders.foreach(_.interrupt()) // if there were still any grpcResponseSenders, register detach time @@ -244,12 +237,26 @@ private[connect] class ExecuteHolder( lastAttachedRpcTimeMs = Some(System.currentTimeMillis()) grpcResponseSenders.clear() } - // remove all cached responses from observer - responseObserver.removeAll() + if (!interrupted) { + cleanup() + } closedTimeMs = Some(System.currentTimeMillis()) } } + /** + * A piece of code that is called only once when this execute holder is closed or the + * interrupted execution thread is terminated. + */ + private[connect] def cleanup(): Unit = { + if (completionCallbackCalled.compareAndSet(false, true)) { + // Remove all cached responses from the observer. + responseObserver.removeAll() + // Post "closed" to UI. + eventsManager.postClosed() + } + } + /** * Spark Connect tags are also added as SparkContext job tags, but to make the tag unique, they * need to be combined with userId and sessionId. diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala index cb0bd8f771ebc..f86298a8b5b98 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala @@ -109,7 +109,7 @@ class SparkConnectServiceE2ESuite extends SparkConnectServerTest { } // Check the async execute cleanup get called Eventually.eventually(timeout(eventuallyTimeout)) { - assert(executeHolder1.completionCallbackCalled) + assert(executeHolder1.completionCallbackCalled.get()) } } } From 5fb0ff9e10b1df266732466790264fd63f159446 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Tue, 24 Sep 2024 21:22:42 -0400 Subject: [PATCH 124/189] [SPARK-49282][CONNECT][SQL] Create a shared SparkSessionBuilder interface ### What changes were proposed in this pull request? This PR adds a shared SparkSessionBuilder interface. It also adds a SparkSessionCompanion interface which is mean should be implemented by all SparkSession companions (a.k.a. `object SparkSession`. This is currently the entry point for session building, in the future we will also add the management of active/default sessions. Finally we add a companion for api.SparkSession. This will bind the implementation that is currently located in `org.apache.spark.sql`. This makes it possible to exclusively work with the interface, instead of selecting an implementation upfront. ### Why are the changes needed? We are creating a shared Scala SQL interface. Building a session is part of this interface. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. I have added tests for the implementation binding. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48229 from hvanhovell/SPARK-49282. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- connector/connect/client/jvm/pom.xml | 7 + .../org/apache/spark/sql/SparkSession.scala | 96 +++------- ...ionBuilderImplementationBindingSuite.scala | 33 ++++ project/MimaExcludes.scala | 46 +++-- .../apache/spark/sql/api/SparkSession.scala | 171 ++++++++++++++++- ...ionBuilderImplementationBindingSuite.scala | 38 ++++ sql/core/pom.xml | 7 + .../org/apache/spark/sql/SparkSession.scala | 179 ++++++------------ ...ionBuilderImplementationBindingSuite.scala | 26 +++ 9 files changed, 388 insertions(+), 215 deletions(-) create mode 100644 connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala create mode 100644 sql/api/src/test/scala/org/apache/spark/sql/api/SparkSessionBuilderImplementationBindingSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala diff --git a/connector/connect/client/jvm/pom.xml b/connector/connect/client/jvm/pom.xml index be358f317481e..e117a0a7451cb 100644 --- a/connector/connect/client/jvm/pom.xml +++ b/connector/connect/client/jvm/pom.xml @@ -88,6 +88,13 @@ scalacheck_${scala.binary.version} test + + org.apache.spark + spark-sql-api_${scala.binary.version} + ${project.version} + tests + test + org.apache.spark spark-common-utils_${scala.binary.version} diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 0663f0186888e..5313369a2c987 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -509,7 +509,7 @@ class SparkSession private[sql] ( // The minimal builder needed to create a spark session. // TODO: implements all methods mentioned in the scaladoc of [[SparkSession]] -object SparkSession extends Logging { +object SparkSession extends api.SparkSessionCompanion with Logging { private val MAX_CACHED_SESSIONS = 100 private val planIdGenerator = new AtomicLong private var server: Option[Process] = None @@ -618,15 +618,15 @@ object SparkSession extends Logging { */ def builder(): Builder = new Builder() - class Builder() extends Logging { + class Builder() extends api.SparkSessionBuilder { // Initialize the connection string of the Spark Connect client builder from SPARK_REMOTE // by default, if it exists. The connection string can be overridden using // the remote() function, as it takes precedence over the SPARK_REMOTE environment variable. private val builder = SparkConnectClient.builder().loadFromEnvironment() private var client: SparkConnectClient = _ - private[this] val options = new scala.collection.mutable.HashMap[String, String] - def remote(connectionString: String): Builder = { + /** @inheritdoc */ + def remote(connectionString: String): this.type = { builder.connectionString(connectionString) this } @@ -638,93 +638,45 @@ object SparkSession extends Logging { * * @since 3.5.0 */ - def interceptor(interceptor: ClientInterceptor): Builder = { + def interceptor(interceptor: ClientInterceptor): this.type = { builder.interceptor(interceptor) this } - private[sql] def client(client: SparkConnectClient): Builder = { + private[sql] def client(client: SparkConnectClient): this.type = { this.client = client this } - /** - * Sets a config option. Options set using this method are automatically propagated to the - * Spark Connect session. Only runtime options are supported. - * - * @since 3.5.0 - */ - def config(key: String, value: String): Builder = synchronized { - options += key -> value - this - } + /** @inheritdoc */ + override def config(key: String, value: String): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to the - * Spark Connect session. Only runtime options are supported. - * - * @since 3.5.0 - */ - def config(key: String, value: Long): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: Long): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to the - * Spark Connect session. Only runtime options are supported. - * - * @since 3.5.0 - */ - def config(key: String, value: Double): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: Double): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to the - * Spark Connect session. Only runtime options are supported. - * - * @since 3.5.0 - */ - def config(key: String, value: Boolean): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: Boolean): this.type = super.config(key, value) - /** - * Sets a config a map of options. Options set using this method are automatically propagated - * to the Spark Connect session. Only runtime options are supported. - * - * @since 3.5.0 - */ - def config(map: Map[String, Any]): Builder = synchronized { - map.foreach { kv: (String, Any) => - { - options += kv._1 -> kv._2.toString - } - } - this - } + /** @inheritdoc */ + override def config(map: Map[String, Any]): this.type = super.config(map) - /** - * Sets a config option. Options set using this method are automatically propagated to both - * `SparkConf` and SparkSession's own configuration. - * - * @since 3.5.0 - */ - def config(map: java.util.Map[String, Any]): Builder = synchronized { - config(map.asScala.toMap) - } + /** @inheritdoc */ + override def config(map: java.util.Map[String, Any]): this.type = super.config(map) + /** @inheritdoc */ @deprecated("enableHiveSupport does not work in Spark Connect") - def enableHiveSupport(): Builder = this + override def enableHiveSupport(): this.type = this + /** @inheritdoc */ @deprecated("master does not work in Spark Connect, please use remote instead") - def master(master: String): Builder = this + override def master(master: String): this.type = this + /** @inheritdoc */ @deprecated("appName does not work in Spark Connect") - def appName(name: String): Builder = this + override def appName(name: String): this.type = this private def tryCreateSessionFromClient(): Option[SparkSession] = { if (client != null && client.isSessionValid) { diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala new file mode 100644 index 0000000000000..ed930882ac2fd --- /dev/null +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql + +import org.apache.spark.sql.api.SparkSessionBuilder +import org.apache.spark.sql.test.{ConnectFunSuite, RemoteSparkSession} + +/** + * Make sure the api.SparkSessionBuilder binds to Connect implementation. + */ +class SparkSessionBuilderImplementationBindingSuite + extends ConnectFunSuite + with api.SparkSessionBuilderImplementationBindingSuite + with RemoteSparkSession { + override protected def configure(builder: SparkSessionBuilder): builder.type = { + // We need to set this configuration because the port used by the server is random. + builder.remote(s"sc://localhost:$serverPort") + } +} diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 972438d0757a7..9a89ebb4797c9 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -125,26 +125,6 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.Observation"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.Observation$"), - // SPARK-49414: Remove Logging from DataFrameReader. - ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.DataFrameReader"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logName"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.log"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logInfo"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logDebug"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logTrace"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logWarning"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logError"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logInfo"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logDebug"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logTrace"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logWarning"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.logError"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.isTraceEnabled"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.initializeLogIfNecessary"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.initializeLogIfNecessary"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.initializeLogIfNecessary$default$2"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.DataFrameReader.initializeForcefully"), - // SPARK-49425: Create a shared DataFrameWriter interface. ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.DataFrameWriter"), @@ -195,7 +175,11 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.SQLImplicits$StringToColumn"), ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SparkSession$implicits$"), ProblemFilters.exclude[InheritedNewAbstractMethodProblem]("org.apache.spark.sql.SQLImplicits.session"), - ) + + // SPARK-49282: Shared SparkSessionBuilder + ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SparkSession$Builder"), + ) ++ loggingExcludes("org.apache.spark.sql.DataFrameReader") ++ + loggingExcludes("org.apache.spark.sql.SparkSession#Builder") // Default exclude rules lazy val defaultExcludes = Seq( @@ -236,6 +220,26 @@ object MimaExcludes { } ) + private def loggingExcludes(fqn: String) = { + Seq( + ProblemFilters.exclude[MissingTypesProblem](fqn), + missingMethod(fqn, "logName"), + missingMethod(fqn, "log"), + missingMethod(fqn, "logInfo"), + missingMethod(fqn, "logDebug"), + missingMethod(fqn, "logTrace"), + missingMethod(fqn, "logWarning"), + missingMethod(fqn, "logError"), + missingMethod(fqn, "isTraceEnabled"), + missingMethod(fqn, "initializeLogIfNecessary"), + missingMethod(fqn, "initializeLogIfNecessary$default$2"), + missingMethod(fqn, "initializeForcefully")) + } + + private def missingMethod(names: String*) = { + ProblemFilters.exclude[DirectMissingMethodProblem](names.mkString(".")) + } + def excludes(version: String): Seq[Problem => Boolean] = version match { case v if v.startsWith("4.0") => v40excludes case _ => Seq() diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 2623db4060ee6..2295c153cd51c 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -25,9 +25,10 @@ import _root_.java.lang import _root_.java.net.URI import _root_.java.util -import org.apache.spark.annotation.{DeveloperApi, Experimental} +import org.apache.spark.annotation.{DeveloperApi, Experimental, Stable} import org.apache.spark.sql.{Encoder, Row, RuntimeConfig} import org.apache.spark.sql.types.StructType +import org.apache.spark.util.SparkClassUtils /** * The entry point to programming Spark with the Dataset and DataFrame API. @@ -541,3 +542,171 @@ abstract class SparkSession extends Serializable with Closeable { */ def stop(): Unit = close() } + +object SparkSession extends SparkSessionCompanion { + private[this] val companion: SparkSessionCompanion = { + val cls = SparkClassUtils.classForName("org.apache.spark.sql.SparkSession") + val mirror = scala.reflect.runtime.currentMirror + val module = mirror.classSymbol(cls).companion.asModule + mirror.reflectModule(module).instance.asInstanceOf[SparkSessionCompanion] + } + + /** @inheritdoc */ + override def builder(): SparkSessionBuilder = companion.builder() +} + +/** + * Companion of a [[SparkSession]]. + */ +private[sql] abstract class SparkSessionCompanion { + + /** + * Creates a [[SparkSessionBuilder]] for constructing a [[SparkSession]]. + * + * @since 2.0.0 + */ + def builder(): SparkSessionBuilder +} + +/** + * Builder for [[SparkSession]]. + */ +@Stable +abstract class SparkSessionBuilder { + protected val options = new scala.collection.mutable.HashMap[String, String] + + /** + * Sets a name for the application, which will be shown in the Spark web UI. If no application + * name is set, a randomly generated name will be used. + * + * @since 2.0.0 + */ + def appName(name: String): this.type = config("spark.app.name", name) + + /** + * Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]" to run + * locally with 4 cores, or "spark://master:7077" to run on a Spark standalone cluster. + * + * @note + * this is only supported in Classic. + * @since 2.0.0 + */ + def master(master: String): this.type = config("spark.master", master) + + /** + * Enables Hive support, including connectivity to a persistent Hive metastore, support for Hive + * serdes, and Hive user-defined functions. + * + * @note + * this is only supported in Classic. + * @since 2.0.0 + */ + def enableHiveSupport(): this.type = config("spark.sql.catalogImplementation", "hive") + + /** + * Sets the Spark Connect remote URL. + * + * @note + * this is only supported in Connect. + * @since 3.5.0 + */ + def remote(connectionString: String): this.type + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @note + * this is only supported in Connect mode. + * @since 2.0.0 + */ + def config(key: String, value: String): this.type = synchronized { + options += key -> value + this + } + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @since 2.0.0 + */ + def config(key: String, value: Long): this.type = synchronized { + options += key -> value.toString + this + } + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @since 2.0.0 + */ + def config(key: String, value: Double): this.type = synchronized { + options += key -> value.toString + this + } + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @since 2.0.0 + */ + def config(key: String, value: Boolean): this.type = synchronized { + options += key -> value.toString + this + } + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @since 3.4.0 + */ + def config(map: Map[String, Any]): this.type = synchronized { + map.foreach { kv: (String, Any) => + { + options += kv._1 -> kv._2.toString + } + } + this + } + + /** + * Sets a config option. Options set using this method are automatically propagated to both + * `SparkConf` and SparkSession's own configuration. + * + * @since 3.4.0 + */ + def config(map: util.Map[String, Any]): this.type = synchronized { + config(map.asScala.toMap) + } + + /** + * Gets an existing [[SparkSession]] or, if there is no existing one, creates a new one based on + * the options set in this builder. + * + * This method first checks whether there is a valid thread-local SparkSession, and if yes, + * return that one. It then checks whether there is a valid global default SparkSession, and if + * yes, return that one. If no valid global default SparkSession exists, the method creates a + * new SparkSession and assigns the newly created SparkSession as the global default. + * + * In case an existing SparkSession is returned, the non-static config options specified in this + * builder will be applied to the existing SparkSession. + * + * @since 2.0.0 + */ + def getOrCreate(): SparkSession + + /** + * Create a new [[SparkSession]]. + * + * This will always return a newly created session. + * + * This method will update the default and/or active session if they are not set. + * + * @since 3.5.0 + */ + def create(): SparkSession +} diff --git a/sql/api/src/test/scala/org/apache/spark/sql/api/SparkSessionBuilderImplementationBindingSuite.scala b/sql/api/src/test/scala/org/apache/spark/sql/api/SparkSessionBuilderImplementationBindingSuite.scala new file mode 100644 index 0000000000000..84b6b85f639a3 --- /dev/null +++ b/sql/api/src/test/scala/org/apache/spark/sql/api/SparkSessionBuilderImplementationBindingSuite.scala @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.api + +// scalastyle:off funsuite +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.functions.sum + +/** + * Test suite for SparkSession implementation binding. + */ +trait SparkSessionBuilderImplementationBindingSuite extends AnyFunSuite with BeforeAndAfterAll { +// scalastyle:on + protected def configure(builder: SparkSessionBuilder): builder.type = builder + + test("range") { + val session = configure(SparkSession.builder()).getOrCreate() + import session.implicits._ + val df = session.range(10).agg(sum("id")).as[Long] + assert(df.head() == 45) + } +} diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 9eb5decb3b515..4352c44a4feda 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -73,6 +73,13 @@ test-jar test + + org.apache.spark + spark-sql-api_${scala.binary.version} + ${project.version} + test-jar + test + org.apache.spark spark-tags_${scala.binary.version} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index 938df206b9792..fe139d629eb24 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -853,129 +853,64 @@ class SparkSession private( @Stable -object SparkSession extends Logging { +object SparkSession extends api.SparkSessionCompanion with Logging { /** * Builder for [[SparkSession]]. */ @Stable - class Builder extends Logging { - - private[this] val options = new scala.collection.mutable.HashMap[String, String] + class Builder extends api.SparkSessionBuilder { private[this] val extensions = new SparkSessionExtensions private[this] var userSuppliedContext: Option[SparkContext] = None - private[spark] def sparkContext(sparkContext: SparkContext): Builder = synchronized { + private[spark] def sparkContext(sparkContext: SparkContext): this.type = synchronized { userSuppliedContext = Option(sparkContext) this } - /** - * Sets a name for the application, which will be shown in the Spark web UI. - * If no application name is set, a randomly generated name will be used. - * - * @since 2.0.0 - */ - def appName(name: String): Builder = config("spark.app.name", name) + /** @inheritdoc */ + override def remote(connectionString: String): this.type = this - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 2.0.0 - */ - def config(key: String, value: String): Builder = synchronized { - options += key -> value - this - } + /** @inheritdoc */ + override def appName(name: String): this.type = super.appName(name) - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 2.0.0 - */ - def config(key: String, value: Long): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: String): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 2.0.0 - */ - def config(key: String, value: Double): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: Long): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 2.0.0 - */ - def config(key: String, value: Boolean): Builder = synchronized { - options += key -> value.toString - this - } + /** @inheritdoc */ + override def config(key: String, value: Double): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 3.4.0 - */ - def config(map: Map[String, Any]): Builder = synchronized { - map.foreach { - kv: (String, Any) => { - options += kv._1 -> kv._2.toString - } - } - this - } + /** @inheritdoc */ + override def config(key: String, value: Boolean): this.type = super.config(key, value) - /** - * Sets a config option. Options set using this method are automatically propagated to - * both `SparkConf` and SparkSession's own configuration. - * - * @since 3.4.0 - */ - def config(map: java.util.Map[String, Any]): Builder = synchronized { - config(map.asScala.toMap) - } + /** @inheritdoc */ + override def config(map: Map[String, Any]): this.type = super.config(map) + + /** @inheritdoc */ + override def config(map: java.util.Map[String, Any]): this.type = super.config(map) /** * Sets a list of config options based on the given `SparkConf`. * * @since 2.0.0 */ - def config(conf: SparkConf): Builder = synchronized { + def config(conf: SparkConf): this.type = synchronized { conf.getAll.foreach { case (k, v) => options += k -> v } this } - /** - * Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]" to - * run locally with 4 cores, or "spark://master:7077" to run on a Spark standalone cluster. - * - * @since 2.0.0 - */ - def master(master: String): Builder = config("spark.master", master) + /** @inheritdoc */ + override def master(master: String): this.type = super.master(master) - /** - * Enables Hive support, including connectivity to a persistent Hive metastore, support for - * Hive serdes, and Hive user-defined functions. - * - * @since 2.0.0 - */ - def enableHiveSupport(): Builder = synchronized { + /** @inheritdoc */ + override def enableHiveSupport(): this.type = synchronized { if (hiveClassesArePresent) { - config(CATALOG_IMPLEMENTATION.key, "hive") + super.enableHiveSupport() } else { throw new IllegalArgumentException( "Unable to instantiate SparkSession with Hive support because " + @@ -989,27 +924,12 @@ object SparkSession extends Logging { * * @since 2.2.0 */ - def withExtensions(f: SparkSessionExtensions => Unit): Builder = synchronized { + def withExtensions(f: SparkSessionExtensions => Unit): this.type = synchronized { f(extensions) this } - /** - * Gets an existing [[SparkSession]] or, if there is no existing one, creates a new - * one based on the options set in this builder. - * - * This method first checks whether there is a valid thread-local SparkSession, - * and if yes, return that one. It then checks whether there is a valid global - * default SparkSession, and if yes, return that one. If no valid global default - * SparkSession exists, the method creates a new SparkSession and assigns the - * newly created SparkSession as the global default. - * - * In case an existing SparkSession is returned, the non-static config options specified in - * this builder will be applied to the existing SparkSession. - * - * @since 2.0.0 - */ - def getOrCreate(): SparkSession = synchronized { + private def build(forceCreate: Boolean): SparkSession = synchronized { val sparkConf = new SparkConf() options.foreach { case (k, v) => sparkConf.set(k, v) } @@ -1017,20 +937,28 @@ object SparkSession extends Logging { assertOnDriver() } + def clearSessionIfDead(session: SparkSession): SparkSession = { + if ((session ne null) && !session.sparkContext.isStopped) { + session + } else { + null + } + } + // Get the session from current thread's active session. - var session = activeThreadSession.get() - if ((session ne null) && !session.sparkContext.isStopped) { - applyModifiableSettings(session, new java.util.HashMap[String, String](options.asJava)) - return session + val active = clearSessionIfDead(activeThreadSession.get()) + if (!forceCreate && (active ne null)) { + applyModifiableSettings(active, new java.util.HashMap[String, String](options.asJava)) + return active } // Global synchronization so we will only set the default session once. SparkSession.synchronized { // If the current thread does not have an active session, get it from the global session. - session = defaultSession.get() - if ((session ne null) && !session.sparkContext.isStopped) { - applyModifiableSettings(session, new java.util.HashMap[String, String](options.asJava)) - return session + val default = clearSessionIfDead(defaultSession.get()) + if (!forceCreate && (default ne null)) { + applyModifiableSettings(default, new java.util.HashMap[String, String](options.asJava)) + return default } // No active nor global default session. Create a new one. @@ -1047,19 +975,28 @@ object SparkSession extends Logging { loadExtensions(extensions) applyExtensions(sparkContext, extensions) - session = new SparkSession(sparkContext, + val session = new SparkSession(sparkContext, existingSharedState = None, parentSessionState = None, extensions, initialSessionOptions = options.toMap, parentManagedJobTags = Map.empty) - setDefaultSession(session) - setActiveSession(session) + if (default eq null) { + setDefaultSession(session) + } + if (active eq null) { + setActiveSession(session) + } registerContextListener(sparkContext) + session } - - return session } + + /** @inheritdoc */ + def getOrCreate(): SparkSession = build(forceCreate = false) + + /** @inheritdoc */ + def create(): SparkSession = build(forceCreate = true) } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala new file mode 100644 index 0000000000000..c4fd16ca5ce59 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderImplementationBindingSuite.scala @@ -0,0 +1,26 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql + +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Make sure the api.SparkSessionBuilder binds to Classic implementation. + */ +class SparkSessionBuilderImplementationBindingSuite + extends SharedSparkSession + with api.SparkSessionBuilderImplementationBindingSuite From 0c234bb1a68c8f419471182d394145c9d48fb3a5 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Tue, 24 Sep 2024 21:24:38 -0400 Subject: [PATCH 125/189] [SPARK-49369][CONNECT][SQL] Add implicit Column conversions ### What changes were proposed in this pull request? This introduces an implicit conversion for the Column companion object that allows a user/developer to create a Column from a catalyst Expression (for Classic) or a proto Expression (Builder) (for Connect). This mostly recreates they had before we refactored the Column API. This comes at the price of adding the an import. ### Why are the changes needed? Improved upgrade experience for Developers and User who create their own Column's from expressions. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? I added it to a couple of places in the code and it works. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48020 from hvanhovell/SPARK-49369. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../sql/connect/ConnectConversions.scala | 40 ++++++++++++++++++- .../scala/org/apache/spark/sql/package.scala | 27 ------------- .../spark/sql/PlanGenerationTestSuite.scala | 1 + .../spark/sql/DataFrameNaFunctions.scala | 17 ++++---- .../scala/org/apache/spark/sql/Dataset.scala | 29 +++++++++----- .../sql/classic/ClassicConversions.scala | 11 ++++- .../sql/internal/RuntimeConfigImpl.scala | 2 +- 7 files changed, 76 insertions(+), 51 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala index 7d81f4ead7857..0344152be86e6 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/ConnectConversions.scala @@ -19,13 +19,15 @@ package org.apache.spark.sql.connect import scala.language.implicitConversions import org.apache.spark.annotation.DeveloperApi +import org.apache.spark.connect.proto import org.apache.spark.sql._ +import org.apache.spark.sql.internal.ProtoColumnNode /** * Conversions from sql interfaces to the Connect specific implementation. * - * This class is mainly used by the implementation. In the case of connect it should be extremely - * rare that a developer needs these classes. + * This class is mainly used by the implementation. It is also meant to be used by extension + * developers. * * We provide both a trait and an object. The trait is useful in situations where an extension * developer needs to use these conversions in a project covering multiple Spark versions. They @@ -46,6 +48,40 @@ trait ConnectConversions { implicit def castToImpl[K, V]( kvds: api.KeyValueGroupedDataset[K, V]): KeyValueGroupedDataset[K, V] = kvds.asInstanceOf[KeyValueGroupedDataset[K, V]] + + /** + * Create a [[Column]] from a [[proto.Expression]] + * + * This method is meant to be used by Connect plugins. We do not guarantee any compatibility + * between (minor) versions. + */ + @DeveloperApi + def column(expr: proto.Expression): Column = { + Column(ProtoColumnNode(expr)) + } + + /** + * Create a [[Column]] using a function that manipulates an [[proto.Expression.Builder]]. + * + * This method is meant to be used by Connect plugins. We do not guarantee any compatibility + * between (minor) versions. + */ + @DeveloperApi + def column(f: proto.Expression.Builder => Unit): Column = { + val builder = proto.Expression.newBuilder() + f(builder) + column(builder.build()) + } + + /** + * Implicit helper that makes it easy to construct a Column from an Expression or an Expression + * builder. This allows developers to create a Column in the same way as in earlier versions of + * Spark (before 4.0). + */ + @DeveloperApi + implicit class ColumnConstructorExt(val c: Column.type) { + def apply(e: proto.Expression): Column = column(e) + } } object ConnectConversions extends ConnectConversions diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala index 154f2b0405fcd..556b472283a37 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala @@ -17,10 +17,7 @@ package org.apache.spark -import org.apache.spark.annotation.DeveloperApi -import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder -import org.apache.spark.sql.internal.ProtoColumnNode package object sql { type DataFrame = Dataset[Row] @@ -28,28 +25,4 @@ package object sql { private[sql] def encoderFor[E: Encoder]: AgnosticEncoder[E] = { implicitly[Encoder[E]].asInstanceOf[AgnosticEncoder[E]] } - - /** - * Create a [[Column]] from a [[proto.Expression]] - * - * This method is meant to be used by Connect plugins. We do not guarantee any compatility - * between (minor) versions. - */ - @DeveloperApi - def column(expr: proto.Expression): Column = { - Column(ProtoColumnNode(expr)) - } - - /** - * Creat a [[Column]] using a function that manipulates an [[proto.Expression.Builder]]. - * - * This method is meant to be used by Connect plugins. We do not guarantee any compatility - * between (minor) versions. - */ - @DeveloperApi - def column(f: proto.Expression.Builder => Unit): Column = { - val builder = proto.Expression.newBuilder() - f(builder) - column(builder.build()) - } } diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala index 315f80e13eff7..c557b54732797 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala @@ -37,6 +37,7 @@ import org.apache.spark.sql.avro.{functions => avroFn} import org.apache.spark.sql.catalyst.ScalaReflection import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.StringEncoder import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.connect.client.SparkConnectClient import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.lit diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala index b356751083fc1..53e12f58edd69 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala @@ -24,7 +24,6 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.functions._ -import org.apache.spark.sql.internal.ExpressionUtils.column import org.apache.spark.sql.types._ /** @@ -122,7 +121,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) (attr.dataType.isInstanceOf[NumericType] && targetColumnType == DoubleType))) { replaceCol(attr, replacementMap) } else { - column(attr) + Column(attr) } } df.select(projections : _*) @@ -131,7 +130,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) protected def fillMap(values: Seq[(String, Any)]): DataFrame = { // Error handling val attrToValue = AttributeMap(values.map { case (colName, replaceValue) => - // Check column name exists + // Check Column name exists val attr = df.resolve(colName) match { case a: Attribute => a case _ => throw QueryExecutionErrors.nestedFieldUnsupportedError(colName) @@ -155,7 +154,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) case v: jl.Integer => fillCol[Integer](attr, v) case v: jl.Boolean => fillCol[Boolean](attr, v.booleanValue()) case v: String => fillCol[String](attr, v) - }.getOrElse(column(attr)) + }.getOrElse(Column(attr)) } df.select(projections : _*) } @@ -165,7 +164,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) * with `replacement`. */ private def fillCol[T](attr: Attribute, replacement: T): Column = { - fillCol(attr.dataType, attr.name, column(attr), replacement) + fillCol(attr.dataType, attr.name, Column(attr), replacement) } /** @@ -192,7 +191,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) val branches = replacementMap.flatMap { case (source, target) => Seq(Literal(source), buildExpr(target)) }.toSeq - column(CaseKeyWhen(attr, branches :+ attr)).as(attr.name) + Column(CaseKeyWhen(attr, branches :+ attr)).as(attr.name) } private def convertToDouble(v: Any): Double = v match { @@ -219,7 +218,7 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) // Filtering condition: // only keep the row if it has at least `minNonNulls` non-null and non-NaN values. val predicate = AtLeastNNonNulls(minNonNulls.getOrElse(cols.size), cols) - df.filter(column(predicate)) + df.filter(Column(predicate)) } private[sql] def fillValue(value: Any, cols: Option[Seq[String]]): DataFrame = { @@ -255,9 +254,9 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) } // Only fill if the column is part of the cols list. if (typeMatches && cols.exists(_.semanticEquals(col))) { - fillCol(col.dataType, col.name, column(col), value) + fillCol(col.dataType, col.name, Column(col), value) } else { - column(col) + Column(col) } } df.select(projections : _*) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 80ec70a7864c3..18fc5787a1583 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -63,7 +63,6 @@ import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, Data import org.apache.spark.sql.execution.python.EvaluatePython import org.apache.spark.sql.execution.stat.StatFunctions import org.apache.spark.sql.internal.{DataFrameWriterImpl, DataFrameWriterV2Impl, MergeIntoWriterImpl, SQLConf} -import org.apache.spark.sql.internal.ExpressionUtils.column import org.apache.spark.sql.internal.TypedAggUtils.withInputType import org.apache.spark.sql.streaming.DataStreamWriter import org.apache.spark.sql.types._ @@ -303,7 +302,7 @@ class Dataset[T] private[sql]( truncate: Int): Seq[Seq[String]] = { val newDf = commandResultOptimized.toDF() val castCols = newDf.logicalPlan.output.map { col => - column(ToPrettyString(col)) + Column(ToPrettyString(col)) } val data = newDf.select(castCols: _*).take(numRows + 1) @@ -505,7 +504,7 @@ class Dataset[T] private[sql]( s"New column names (${colNames.size}): " + colNames.mkString(", ")) val newCols = logicalPlan.output.zip(colNames).map { case (oldAttribute, newName) => - column(oldAttribute).as(newName) + Column(oldAttribute).as(newName) } select(newCols : _*) } @@ -760,18 +759,18 @@ class Dataset[T] private[sql]( /** @inheritdoc */ def col(colName: String): Column = colName match { case "*" => - column(ResolvedStar(queryExecution.analyzed.output)) + Column(ResolvedStar(queryExecution.analyzed.output)) case _ => if (sparkSession.sessionState.conf.supportQuotedRegexColumnName) { colRegex(colName) } else { - column(addDataFrameIdToCol(resolve(colName))) + Column(addDataFrameIdToCol(resolve(colName))) } } /** @inheritdoc */ def metadataColumn(colName: String): Column = - column(queryExecution.analyzed.getMetadataAttributeByName(colName)) + Column(queryExecution.analyzed.getMetadataAttributeByName(colName)) // Attach the dataset id and column position to the column reference, so that we can detect // ambiguous self-join correctly. See the rule `DetectAmbiguousSelfJoin`. @@ -797,11 +796,11 @@ class Dataset[T] private[sql]( val caseSensitive = sparkSession.sessionState.conf.caseSensitiveAnalysis colName match { case ParserUtils.escapedIdentifier(columnNameRegex) => - column(UnresolvedRegex(columnNameRegex, None, caseSensitive)) + Column(UnresolvedRegex(columnNameRegex, None, caseSensitive)) case ParserUtils.qualifiedEscapedIdentifier(nameParts, columnNameRegex) => - column(UnresolvedRegex(columnNameRegex, Some(nameParts), caseSensitive)) + Column(UnresolvedRegex(columnNameRegex, Some(nameParts), caseSensitive)) case _ => - column(addDataFrameIdToCol(resolve(colName))) + Column(addDataFrameIdToCol(resolve(colName))) } } @@ -1194,7 +1193,7 @@ class Dataset[T] private[sql]( resolver(field.name, colName) } match { case Some((colName: String, col: Column)) => col.as(colName) - case _ => column(field) + case _ => Column(field) } } @@ -1264,7 +1263,7 @@ class Dataset[T] private[sql]( val allColumns = queryExecution.analyzed.output val remainingCols = allColumns.filter { attribute => colNames.forall(n => !resolver(attribute.name, n)) - }.map(attribute => column(attribute)) + }.map(attribute => Column(attribute)) if (remainingCols.size == allColumns.size) { toDF() } else { @@ -1975,6 +1974,14 @@ class Dataset[T] private[sql]( // For Python API //////////////////////////////////////////////////////////////////////////// + /** + * It adds a new long column with the name `name` that increases one by one. + * This is for 'distributed-sequence' default index in pandas API on Spark. + */ + private[sql] def withSequenceColumn(name: String) = { + select(Column(DistributedSequenceID()).alias(name), col("*")) + } + /** * Converts a JavaRDD to a PythonRDD. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala index af91b57a6848b..8c3223fa72f55 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/ClassicConversions.scala @@ -20,11 +20,13 @@ import scala.language.implicitConversions import org.apache.spark.annotation.DeveloperApi import org.apache.spark.sql._ +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.internal.ExpressionUtils /** * Conversions from sql interfaces to the Classic specific implementation. * - * This class is mainly used by the implementation, but is also meant to be used by extension + * This class is mainly used by the implementation. It is also meant to be used by extension * developers. * * We provide both a trait and an object. The trait is useful in situations where an extension @@ -45,6 +47,13 @@ trait ClassicConversions { implicit def castToImpl[K, V](kvds: api.KeyValueGroupedDataset[K, V]) : KeyValueGroupedDataset[K, V] = kvds.asInstanceOf[KeyValueGroupedDataset[K, V]] + + /** + * Helper that makes it easy to construct a Column from an Expression. + */ + implicit class ColumnConstructorExt(val c: Column.type) { + def apply(e: Expression): Column = ExpressionUtils.column(e) + } } object ClassicConversions extends ClassicConversions diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala index ca439cdb89958..f25ca387db299 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/RuntimeConfigImpl.scala @@ -84,7 +84,7 @@ class RuntimeConfigImpl private[sql](val sqlConf: SQLConf = new SQLConf) extends sqlConf.contains(key) } - private def requireNonStaticConf(key: String): Unit = { + private[sql] def requireNonStaticConf(key: String): Unit = { if (SQLConf.isStaticConfigKey(key)) { throw QueryCompilationErrors.cannotModifyValueOfStaticConfigError(key) } From 828b1f94734af8a629e80b1ec2d7f25326c69411 Mon Sep 17 00:00:00 2001 From: bogao007 Date: Wed, 25 Sep 2024 11:05:15 +0900 Subject: [PATCH 126/189] [SPARK-49463] Support ListState for TransformWithStateInPandas ### What changes were proposed in this pull request? Support ListState for TransformWithStateInPandas ### Why are the changes needed? Adding new functionality for TransformWithStateInPandas ### Does this PR introduce _any_ user-facing change? Yes ### How was this patch tested? Added new unit tests. ### Was this patch authored or co-authored using generative AI tooling? No Closes #47933 from bogao007/list-state. Authored-by: bogao007 Signed-off-by: Jungtaek Lim --- python/pyspark/sql/pandas/types.py | 36 + .../pyspark/sql/streaming/StateMessage_pb2.py | 71 +- .../sql/streaming/StateMessage_pb2.pyi | 313 +- .../sql/streaming/list_state_client.py | 187 + .../sql/streaming/stateful_processor.py | 87 +- .../stateful_processor_api_client.py | 45 +- .../sql/streaming/value_state_client.py | 8 +- .../test_pandas_transform_with_state.py | 63 + .../apache/spark/sql/internal/SQLConf.scala | 11 + .../execution/streaming/StateMessage.proto | 27 + .../streaming/state/StateMessage.java | 4942 +++++++++++++++-- ...ansformWithStateInPandasDeserializer.scala | 60 + ...ansformWithStateInPandasPythonRunner.scala | 3 +- ...ransformWithStateInPandasStateServer.scala | 205 +- ...ormWithStateInPandasStateServerSuite.scala | 157 +- 15 files changed, 5570 insertions(+), 645 deletions(-) create mode 100644 python/pyspark/sql/streaming/list_state_client.py create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasDeserializer.scala diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 53c72304adfaa..57e46901013fe 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -53,12 +53,17 @@ ) from pyspark.errors import PySparkTypeError, UnsupportedOperationException, PySparkValueError from pyspark.loose_version import LooseVersion +from pyspark.sql.utils import has_numpy + +if has_numpy: + import numpy as np if TYPE_CHECKING: import pandas as pd import pyarrow as pa from pyspark.sql.pandas._typing import SeriesLike as PandasSeriesLike + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike def to_arrow_type( @@ -1344,3 +1349,34 @@ def _deduplicate_field_names(dt: DataType) -> DataType: ) else: return dt + + +def _to_numpy_type(type: DataType) -> Optional["np.dtype"]: + """Convert Spark data type to NumPy type.""" + import numpy as np + + if type == ByteType(): + return np.dtype("int8") + elif type == ShortType(): + return np.dtype("int16") + elif type == IntegerType(): + return np.dtype("int32") + elif type == LongType(): + return np.dtype("int64") + elif type == FloatType(): + return np.dtype("float32") + elif type == DoubleType(): + return np.dtype("float64") + return None + + +def convert_pandas_using_numpy_type( + df: "PandasDataFrameLike", schema: StructType +) -> "PandasDataFrameLike": + for field in schema.fields: + if isinstance( + field.dataType, (ByteType, ShortType, LongType, FloatType, DoubleType, IntegerType) + ): + np_type = _to_numpy_type(field.dataType) + df[field.name] = df[field.name].astype(np_type) + return df diff --git a/python/pyspark/sql/streaming/StateMessage_pb2.py b/python/pyspark/sql/streaming/StateMessage_pb2.py index a22f004fd3048..e75d0394ea0f5 100644 --- a/python/pyspark/sql/streaming/StateMessage_pb2.py +++ b/python/pyspark/sql/streaming/StateMessage_pb2.py @@ -16,12 +16,14 @@ # # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE # source: StateMessage.proto +# Protobuf Python Version: 5.27.3 """Generated protocol buffer code.""" -from google.protobuf.internal import builder as _builder from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) @@ -29,45 +31,54 @@ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x12StateMessage.proto\x12.org.apache.spark.sql.execution.streaming.state"\xe9\x02\n\x0cStateRequest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x66\n\x15statefulProcessorCall\x18\x02 \x01(\x0b\x32\x45.org.apache.spark.sql.execution.streaming.state.StatefulProcessorCallH\x00\x12\x64\n\x14stateVariableRequest\x18\x03 \x01(\x0b\x32\x44.org.apache.spark.sql.execution.streaming.state.StateVariableRequestH\x00\x12p\n\x1aimplicitGroupingKeyRequest\x18\x04 \x01(\x0b\x32J.org.apache.spark.sql.execution.streaming.state.ImplicitGroupingKeyRequestH\x00\x42\x08\n\x06method"H\n\rStateResponse\x12\x12\n\nstatusCode\x18\x01 \x01(\x05\x12\x14\n\x0c\x65rrorMessage\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x0c"\x89\x03\n\x15StatefulProcessorCall\x12X\n\x0esetHandleState\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.SetHandleStateH\x00\x12Y\n\rgetValueState\x18\x02 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x12X\n\x0cgetListState\x18\x03 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x12W\n\x0bgetMapState\x18\x04 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x42\x08\n\x06method"z\n\x14StateVariableRequest\x12X\n\x0evalueStateCall\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.ValueStateCallH\x00\x42\x08\n\x06method"\xe0\x01\n\x1aImplicitGroupingKeyRequest\x12X\n\x0esetImplicitKey\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.SetImplicitKeyH\x00\x12^\n\x11removeImplicitKey\x18\x02 \x01(\x0b\x32\x41.org.apache.spark.sql.execution.streaming.state.RemoveImplicitKeyH\x00\x42\x08\n\x06method"}\n\x10StateCallCommand\x12\x11\n\tstateName\x18\x01 \x01(\t\x12\x0e\n\x06schema\x18\x02 \x01(\t\x12\x46\n\x03ttl\x18\x03 \x01(\x0b\x32\x39.org.apache.spark.sql.execution.streaming.state.TTLConfig"\xe1\x02\n\x0eValueStateCall\x12\x11\n\tstateName\x18\x01 \x01(\t\x12H\n\x06\x65xists\x18\x02 \x01(\x0b\x32\x36.org.apache.spark.sql.execution.streaming.state.ExistsH\x00\x12\x42\n\x03get\x18\x03 \x01(\x0b\x32\x33.org.apache.spark.sql.execution.streaming.state.GetH\x00\x12\\\n\x10valueStateUpdate\x18\x04 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.ValueStateUpdateH\x00\x12\x46\n\x05\x63lear\x18\x05 \x01(\x0b\x32\x35.org.apache.spark.sql.execution.streaming.state.ClearH\x00\x42\x08\n\x06method"\x1d\n\x0eSetImplicitKey\x12\x0b\n\x03key\x18\x01 \x01(\x0c"\x13\n\x11RemoveImplicitKey"\x08\n\x06\x45xists"\x05\n\x03Get"!\n\x10ValueStateUpdate\x12\r\n\x05value\x18\x01 \x01(\x0c"\x07\n\x05\x43lear"\\\n\x0eSetHandleState\x12J\n\x05state\x18\x01 \x01(\x0e\x32;.org.apache.spark.sql.execution.streaming.state.HandleState"\x1f\n\tTTLConfig\x12\x12\n\ndurationMs\x18\x01 \x01(\x05*K\n\x0bHandleState\x12\x0b\n\x07\x43REATED\x10\x00\x12\x0f\n\x0bINITIALIZED\x10\x01\x12\x12\n\x0e\x44\x41TA_PROCESSED\x10\x02\x12\n\n\x06\x43LOSED\x10\x03\x62\x06proto3' # noqa: E501 + b'\n\x12StateMessage.proto\x12.org.apache.spark.sql.execution.streaming.state"\xe9\x02\n\x0cStateRequest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x66\n\x15statefulProcessorCall\x18\x02 \x01(\x0b\x32\x45.org.apache.spark.sql.execution.streaming.state.StatefulProcessorCallH\x00\x12\x64\n\x14stateVariableRequest\x18\x03 \x01(\x0b\x32\x44.org.apache.spark.sql.execution.streaming.state.StateVariableRequestH\x00\x12p\n\x1aimplicitGroupingKeyRequest\x18\x04 \x01(\x0b\x32J.org.apache.spark.sql.execution.streaming.state.ImplicitGroupingKeyRequestH\x00\x42\x08\n\x06method"H\n\rStateResponse\x12\x12\n\nstatusCode\x18\x01 \x01(\x05\x12\x14\n\x0c\x65rrorMessage\x18\x02 \x01(\t\x12\r\n\x05value\x18\x03 \x01(\x0c"\x89\x03\n\x15StatefulProcessorCall\x12X\n\x0esetHandleState\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.SetHandleStateH\x00\x12Y\n\rgetValueState\x18\x02 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x12X\n\x0cgetListState\x18\x03 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x12W\n\x0bgetMapState\x18\x04 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.StateCallCommandH\x00\x42\x08\n\x06method"\xd2\x01\n\x14StateVariableRequest\x12X\n\x0evalueStateCall\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.ValueStateCallH\x00\x12V\n\rlistStateCall\x18\x02 \x01(\x0b\x32=.org.apache.spark.sql.execution.streaming.state.ListStateCallH\x00\x42\x08\n\x06method"\xe0\x01\n\x1aImplicitGroupingKeyRequest\x12X\n\x0esetImplicitKey\x18\x01 \x01(\x0b\x32>.org.apache.spark.sql.execution.streaming.state.SetImplicitKeyH\x00\x12^\n\x11removeImplicitKey\x18\x02 \x01(\x0b\x32\x41.org.apache.spark.sql.execution.streaming.state.RemoveImplicitKeyH\x00\x42\x08\n\x06method"}\n\x10StateCallCommand\x12\x11\n\tstateName\x18\x01 \x01(\t\x12\x0e\n\x06schema\x18\x02 \x01(\t\x12\x46\n\x03ttl\x18\x03 \x01(\x0b\x32\x39.org.apache.spark.sql.execution.streaming.state.TTLConfig"\xe1\x02\n\x0eValueStateCall\x12\x11\n\tstateName\x18\x01 \x01(\t\x12H\n\x06\x65xists\x18\x02 \x01(\x0b\x32\x36.org.apache.spark.sql.execution.streaming.state.ExistsH\x00\x12\x42\n\x03get\x18\x03 \x01(\x0b\x32\x33.org.apache.spark.sql.execution.streaming.state.GetH\x00\x12\\\n\x10valueStateUpdate\x18\x04 \x01(\x0b\x32@.org.apache.spark.sql.execution.streaming.state.ValueStateUpdateH\x00\x12\x46\n\x05\x63lear\x18\x05 \x01(\x0b\x32\x35.org.apache.spark.sql.execution.streaming.state.ClearH\x00\x42\x08\n\x06method"\x90\x04\n\rListStateCall\x12\x11\n\tstateName\x18\x01 \x01(\t\x12H\n\x06\x65xists\x18\x02 \x01(\x0b\x32\x36.org.apache.spark.sql.execution.streaming.state.ExistsH\x00\x12T\n\x0clistStateGet\x18\x03 \x01(\x0b\x32<.org.apache.spark.sql.execution.streaming.state.ListStateGetH\x00\x12T\n\x0clistStatePut\x18\x04 \x01(\x0b\x32<.org.apache.spark.sql.execution.streaming.state.ListStatePutH\x00\x12R\n\x0b\x61ppendValue\x18\x05 \x01(\x0b\x32;.org.apache.spark.sql.execution.streaming.state.AppendValueH\x00\x12P\n\nappendList\x18\x06 \x01(\x0b\x32:.org.apache.spark.sql.execution.streaming.state.AppendListH\x00\x12\x46\n\x05\x63lear\x18\x07 \x01(\x0b\x32\x35.org.apache.spark.sql.execution.streaming.state.ClearH\x00\x42\x08\n\x06method"\x1d\n\x0eSetImplicitKey\x12\x0b\n\x03key\x18\x01 \x01(\x0c"\x13\n\x11RemoveImplicitKey"\x08\n\x06\x45xists"\x05\n\x03Get"!\n\x10ValueStateUpdate\x12\r\n\x05value\x18\x01 \x01(\x0c"\x07\n\x05\x43lear""\n\x0cListStateGet\x12\x12\n\niteratorId\x18\x01 \x01(\t"\x0e\n\x0cListStatePut"\x1c\n\x0b\x41ppendValue\x12\r\n\x05value\x18\x01 \x01(\x0c"\x0c\n\nAppendList"\\\n\x0eSetHandleState\x12J\n\x05state\x18\x01 \x01(\x0e\x32;.org.apache.spark.sql.execution.streaming.state.HandleState"\x1f\n\tTTLConfig\x12\x12\n\ndurationMs\x18\x01 \x01(\x05*K\n\x0bHandleState\x12\x0b\n\x07\x43REATED\x10\x00\x12\x0f\n\x0bINITIALIZED\x10\x01\x12\x12\n\x0e\x44\x41TA_PROCESSED\x10\x02\x12\n\n\x06\x43LOSED\x10\x03\x62\x06proto3' # noqa: E501 ) _globals = globals() - _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "StateMessage_pb2", _globals) if not _descriptor._USE_C_DESCRIPTORS: - DESCRIPTOR._options = None - _globals["_HANDLESTATE"]._serialized_start = 1978 - _globals["_HANDLESTATE"]._serialized_end = 2053 + DESCRIPTOR._loaded_options = None + _globals["_HANDLESTATE"]._serialized_start = 2694 + _globals["_HANDLESTATE"]._serialized_end = 2769 _globals["_STATEREQUEST"]._serialized_start = 71 _globals["_STATEREQUEST"]._serialized_end = 432 _globals["_STATERESPONSE"]._serialized_start = 434 _globals["_STATERESPONSE"]._serialized_end = 506 _globals["_STATEFULPROCESSORCALL"]._serialized_start = 509 _globals["_STATEFULPROCESSORCALL"]._serialized_end = 902 - _globals["_STATEVARIABLEREQUEST"]._serialized_start = 904 - _globals["_STATEVARIABLEREQUEST"]._serialized_end = 1026 - _globals["_IMPLICITGROUPINGKEYREQUEST"]._serialized_start = 1029 - _globals["_IMPLICITGROUPINGKEYREQUEST"]._serialized_end = 1253 - _globals["_STATECALLCOMMAND"]._serialized_start = 1255 - _globals["_STATECALLCOMMAND"]._serialized_end = 1380 - _globals["_VALUESTATECALL"]._serialized_start = 1383 - _globals["_VALUESTATECALL"]._serialized_end = 1736 - _globals["_SETIMPLICITKEY"]._serialized_start = 1738 - _globals["_SETIMPLICITKEY"]._serialized_end = 1767 - _globals["_REMOVEIMPLICITKEY"]._serialized_start = 1769 - _globals["_REMOVEIMPLICITKEY"]._serialized_end = 1788 - _globals["_EXISTS"]._serialized_start = 1790 - _globals["_EXISTS"]._serialized_end = 1798 - _globals["_GET"]._serialized_start = 1800 - _globals["_GET"]._serialized_end = 1805 - _globals["_VALUESTATEUPDATE"]._serialized_start = 1807 - _globals["_VALUESTATEUPDATE"]._serialized_end = 1840 - _globals["_CLEAR"]._serialized_start = 1842 - _globals["_CLEAR"]._serialized_end = 1849 - _globals["_SETHANDLESTATE"]._serialized_start = 1851 - _globals["_SETHANDLESTATE"]._serialized_end = 1943 - _globals["_TTLCONFIG"]._serialized_start = 1945 - _globals["_TTLCONFIG"]._serialized_end = 1976 + _globals["_STATEVARIABLEREQUEST"]._serialized_start = 905 + _globals["_STATEVARIABLEREQUEST"]._serialized_end = 1115 + _globals["_IMPLICITGROUPINGKEYREQUEST"]._serialized_start = 1118 + _globals["_IMPLICITGROUPINGKEYREQUEST"]._serialized_end = 1342 + _globals["_STATECALLCOMMAND"]._serialized_start = 1344 + _globals["_STATECALLCOMMAND"]._serialized_end = 1469 + _globals["_VALUESTATECALL"]._serialized_start = 1472 + _globals["_VALUESTATECALL"]._serialized_end = 1825 + _globals["_LISTSTATECALL"]._serialized_start = 1828 + _globals["_LISTSTATECALL"]._serialized_end = 2356 + _globals["_SETIMPLICITKEY"]._serialized_start = 2358 + _globals["_SETIMPLICITKEY"]._serialized_end = 2387 + _globals["_REMOVEIMPLICITKEY"]._serialized_start = 2389 + _globals["_REMOVEIMPLICITKEY"]._serialized_end = 2408 + _globals["_EXISTS"]._serialized_start = 2410 + _globals["_EXISTS"]._serialized_end = 2418 + _globals["_GET"]._serialized_start = 2420 + _globals["_GET"]._serialized_end = 2425 + _globals["_VALUESTATEUPDATE"]._serialized_start = 2427 + _globals["_VALUESTATEUPDATE"]._serialized_end = 2460 + _globals["_CLEAR"]._serialized_start = 2462 + _globals["_CLEAR"]._serialized_end = 2469 + _globals["_LISTSTATEGET"]._serialized_start = 2471 + _globals["_LISTSTATEGET"]._serialized_end = 2505 + _globals["_LISTSTATEPUT"]._serialized_start = 2507 + _globals["_LISTSTATEPUT"]._serialized_end = 2521 + _globals["_APPENDVALUE"]._serialized_start = 2523 + _globals["_APPENDVALUE"]._serialized_end = 2551 + _globals["_APPENDLIST"]._serialized_start = 2553 + _globals["_APPENDLIST"]._serialized_end = 2565 + _globals["_SETHANDLESTATE"]._serialized_start = 2567 + _globals["_SETHANDLESTATE"]._serialized_end = 2659 + _globals["_TTLCONFIG"]._serialized_start = 2661 + _globals["_TTLCONFIG"]._serialized_end = 2692 # @@protoc_insertion_point(module_scope) diff --git a/python/pyspark/sql/streaming/StateMessage_pb2.pyi b/python/pyspark/sql/streaming/StateMessage_pb2.pyi index 1ab48a27c8f87..b1f5f0f7d2a1e 100644 --- a/python/pyspark/sql/streaming/StateMessage_pb2.pyi +++ b/python/pyspark/sql/streaming/StateMessage_pb2.pyi @@ -13,167 +13,238 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# + from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper from google.protobuf import descriptor as _descriptor from google.protobuf import message as _message -from typing import ClassVar, Mapping, Optional, Union +from typing import ( + ClassVar as _ClassVar, + Mapping as _Mapping, + Optional as _Optional, + Union as _Union, +) -CLOSED: HandleState -CREATED: HandleState -DATA_PROCESSED: HandleState DESCRIPTOR: _descriptor.FileDescriptor -INITIALIZED: HandleState -class Clear(_message.Message): - __slots__ = () - def __init__(self) -> None: ... - -class Exists(_message.Message): - __slots__ = () - def __init__(self) -> None: ... - -class Get(_message.Message): - __slots__ = () - def __init__(self) -> None: ... - -class ImplicitGroupingKeyRequest(_message.Message): - __slots__ = ["removeImplicitKey", "setImplicitKey"] - REMOVEIMPLICITKEY_FIELD_NUMBER: ClassVar[int] - SETIMPLICITKEY_FIELD_NUMBER: ClassVar[int] - removeImplicitKey: RemoveImplicitKey - setImplicitKey: SetImplicitKey - def __init__( - self, - setImplicitKey: Optional[Union[SetImplicitKey, Mapping]] = ..., - removeImplicitKey: Optional[Union[RemoveImplicitKey, Mapping]] = ..., - ) -> None: ... - -class RemoveImplicitKey(_message.Message): +class HandleState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): __slots__ = () - def __init__(self) -> None: ... - -class SetHandleState(_message.Message): - __slots__ = ["state"] - STATE_FIELD_NUMBER: ClassVar[int] - state: HandleState - def __init__(self, state: Optional[Union[HandleState, str]] = ...) -> None: ... - -class SetImplicitKey(_message.Message): - __slots__ = ["key"] - KEY_FIELD_NUMBER: ClassVar[int] - key: bytes - def __init__(self, key: Optional[bytes] = ...) -> None: ... + CREATED: _ClassVar[HandleState] + INITIALIZED: _ClassVar[HandleState] + DATA_PROCESSED: _ClassVar[HandleState] + CLOSED: _ClassVar[HandleState] -class StateCallCommand(_message.Message): - __slots__ = ["schema", "stateName", "ttl"] - SCHEMA_FIELD_NUMBER: ClassVar[int] - STATENAME_FIELD_NUMBER: ClassVar[int] - TTL_FIELD_NUMBER: ClassVar[int] - schema: str - stateName: str - ttl: TTLConfig - def __init__( - self, - stateName: Optional[str] = ..., - schema: Optional[str] = ..., - ttl: Optional[Union[TTLConfig, Mapping]] = ..., - ) -> None: ... +CREATED: HandleState +INITIALIZED: HandleState +DATA_PROCESSED: HandleState +CLOSED: HandleState class StateRequest(_message.Message): - __slots__ = [ - "implicitGroupingKeyRequest", - "stateVariableRequest", - "statefulProcessorCall", + __slots__ = ( "version", - ] - IMPLICITGROUPINGKEYREQUEST_FIELD_NUMBER: ClassVar[int] - STATEFULPROCESSORCALL_FIELD_NUMBER: ClassVar[int] - STATEVARIABLEREQUEST_FIELD_NUMBER: ClassVar[int] - VERSION_FIELD_NUMBER: ClassVar[int] - implicitGroupingKeyRequest: ImplicitGroupingKeyRequest - stateVariableRequest: StateVariableRequest - statefulProcessorCall: StatefulProcessorCall + "statefulProcessorCall", + "stateVariableRequest", + "implicitGroupingKeyRequest", + ) + VERSION_FIELD_NUMBER: _ClassVar[int] + STATEFULPROCESSORCALL_FIELD_NUMBER: _ClassVar[int] + STATEVARIABLEREQUEST_FIELD_NUMBER: _ClassVar[int] + IMPLICITGROUPINGKEYREQUEST_FIELD_NUMBER: _ClassVar[int] version: int + statefulProcessorCall: StatefulProcessorCall + stateVariableRequest: StateVariableRequest + implicitGroupingKeyRequest: ImplicitGroupingKeyRequest def __init__( self, - version: Optional[int] = ..., - statefulProcessorCall: Optional[Union[StatefulProcessorCall, Mapping]] = ..., - stateVariableRequest: Optional[Union[StateVariableRequest, Mapping]] = ..., - implicitGroupingKeyRequest: Optional[Union[ImplicitGroupingKeyRequest, Mapping]] = ..., + version: _Optional[int] = ..., + statefulProcessorCall: _Optional[_Union[StatefulProcessorCall, _Mapping]] = ..., + stateVariableRequest: _Optional[_Union[StateVariableRequest, _Mapping]] = ..., + implicitGroupingKeyRequest: _Optional[_Union[ImplicitGroupingKeyRequest, _Mapping]] = ..., ) -> None: ... class StateResponse(_message.Message): - __slots__ = ["errorMessage", "statusCode", "value"] - ERRORMESSAGE_FIELD_NUMBER: ClassVar[int] - STATUSCODE_FIELD_NUMBER: ClassVar[int] - VALUE_FIELD_NUMBER: ClassVar[int] - errorMessage: str + __slots__ = ("statusCode", "errorMessage", "value") + STATUSCODE_FIELD_NUMBER: _ClassVar[int] + ERRORMESSAGE_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] statusCode: int + errorMessage: str value: bytes def __init__( self, - statusCode: Optional[int] = ..., - errorMessage: Optional[str] = ..., - value: Optional[bytes] = ..., + statusCode: _Optional[int] = ..., + errorMessage: _Optional[str] = ..., + value: _Optional[bytes] = ..., ) -> None: ... -class StateVariableRequest(_message.Message): - __slots__ = ["valueStateCall"] - VALUESTATECALL_FIELD_NUMBER: ClassVar[int] - valueStateCall: ValueStateCall - def __init__(self, valueStateCall: Optional[Union[ValueStateCall, Mapping]] = ...) -> None: ... - class StatefulProcessorCall(_message.Message): - __slots__ = ["getListState", "getMapState", "getValueState", "setHandleState"] - GETLISTSTATE_FIELD_NUMBER: ClassVar[int] - GETMAPSTATE_FIELD_NUMBER: ClassVar[int] - GETVALUESTATE_FIELD_NUMBER: ClassVar[int] - SETHANDLESTATE_FIELD_NUMBER: ClassVar[int] + __slots__ = ("setHandleState", "getValueState", "getListState", "getMapState") + SETHANDLESTATE_FIELD_NUMBER: _ClassVar[int] + GETVALUESTATE_FIELD_NUMBER: _ClassVar[int] + GETLISTSTATE_FIELD_NUMBER: _ClassVar[int] + GETMAPSTATE_FIELD_NUMBER: _ClassVar[int] + setHandleState: SetHandleState + getValueState: StateCallCommand getListState: StateCallCommand getMapState: StateCallCommand - getValueState: StateCallCommand - setHandleState: SetHandleState def __init__( self, - setHandleState: Optional[Union[SetHandleState, Mapping]] = ..., - getValueState: Optional[Union[StateCallCommand, Mapping]] = ..., - getListState: Optional[Union[StateCallCommand, Mapping]] = ..., - getMapState: Optional[Union[StateCallCommand, Mapping]] = ..., + setHandleState: _Optional[_Union[SetHandleState, _Mapping]] = ..., + getValueState: _Optional[_Union[StateCallCommand, _Mapping]] = ..., + getListState: _Optional[_Union[StateCallCommand, _Mapping]] = ..., + getMapState: _Optional[_Union[StateCallCommand, _Mapping]] = ..., ) -> None: ... -class TTLConfig(_message.Message): - __slots__ = ["durationMs"] - DURATIONMS_FIELD_NUMBER: ClassVar[int] - durationMs: int - def __init__(self, durationMs: Optional[int] = ...) -> None: ... +class StateVariableRequest(_message.Message): + __slots__ = ("valueStateCall", "listStateCall") + VALUESTATECALL_FIELD_NUMBER: _ClassVar[int] + LISTSTATECALL_FIELD_NUMBER: _ClassVar[int] + valueStateCall: ValueStateCall + listStateCall: ListStateCall + def __init__( + self, + valueStateCall: _Optional[_Union[ValueStateCall, _Mapping]] = ..., + listStateCall: _Optional[_Union[ListStateCall, _Mapping]] = ..., + ) -> None: ... + +class ImplicitGroupingKeyRequest(_message.Message): + __slots__ = ("setImplicitKey", "removeImplicitKey") + SETIMPLICITKEY_FIELD_NUMBER: _ClassVar[int] + REMOVEIMPLICITKEY_FIELD_NUMBER: _ClassVar[int] + setImplicitKey: SetImplicitKey + removeImplicitKey: RemoveImplicitKey + def __init__( + self, + setImplicitKey: _Optional[_Union[SetImplicitKey, _Mapping]] = ..., + removeImplicitKey: _Optional[_Union[RemoveImplicitKey, _Mapping]] = ..., + ) -> None: ... + +class StateCallCommand(_message.Message): + __slots__ = ("stateName", "schema", "ttl") + STATENAME_FIELD_NUMBER: _ClassVar[int] + SCHEMA_FIELD_NUMBER: _ClassVar[int] + TTL_FIELD_NUMBER: _ClassVar[int] + stateName: str + schema: str + ttl: TTLConfig + def __init__( + self, + stateName: _Optional[str] = ..., + schema: _Optional[str] = ..., + ttl: _Optional[_Union[TTLConfig, _Mapping]] = ..., + ) -> None: ... class ValueStateCall(_message.Message): - __slots__ = ["clear", "exists", "get", "stateName", "valueStateUpdate"] - CLEAR_FIELD_NUMBER: ClassVar[int] - EXISTS_FIELD_NUMBER: ClassVar[int] - GET_FIELD_NUMBER: ClassVar[int] - STATENAME_FIELD_NUMBER: ClassVar[int] - VALUESTATEUPDATE_FIELD_NUMBER: ClassVar[int] - clear: Clear + __slots__ = ("stateName", "exists", "get", "valueStateUpdate", "clear") + STATENAME_FIELD_NUMBER: _ClassVar[int] + EXISTS_FIELD_NUMBER: _ClassVar[int] + GET_FIELD_NUMBER: _ClassVar[int] + VALUESTATEUPDATE_FIELD_NUMBER: _ClassVar[int] + CLEAR_FIELD_NUMBER: _ClassVar[int] + stateName: str exists: Exists get: Get - stateName: str valueStateUpdate: ValueStateUpdate + clear: Clear + def __init__( + self, + stateName: _Optional[str] = ..., + exists: _Optional[_Union[Exists, _Mapping]] = ..., + get: _Optional[_Union[Get, _Mapping]] = ..., + valueStateUpdate: _Optional[_Union[ValueStateUpdate, _Mapping]] = ..., + clear: _Optional[_Union[Clear, _Mapping]] = ..., + ) -> None: ... + +class ListStateCall(_message.Message): + __slots__ = ( + "stateName", + "exists", + "listStateGet", + "listStatePut", + "appendValue", + "appendList", + "clear", + ) + STATENAME_FIELD_NUMBER: _ClassVar[int] + EXISTS_FIELD_NUMBER: _ClassVar[int] + LISTSTATEGET_FIELD_NUMBER: _ClassVar[int] + LISTSTATEPUT_FIELD_NUMBER: _ClassVar[int] + APPENDVALUE_FIELD_NUMBER: _ClassVar[int] + APPENDLIST_FIELD_NUMBER: _ClassVar[int] + CLEAR_FIELD_NUMBER: _ClassVar[int] + stateName: str + exists: Exists + listStateGet: ListStateGet + listStatePut: ListStatePut + appendValue: AppendValue + appendList: AppendList + clear: Clear def __init__( self, - stateName: Optional[str] = ..., - exists: Optional[Union[Exists, Mapping]] = ..., - get: Optional[Union[Get, Mapping]] = ..., - valueStateUpdate: Optional[Union[ValueStateUpdate, Mapping]] = ..., - clear: Optional[Union[Clear, Mapping]] = ..., + stateName: _Optional[str] = ..., + exists: _Optional[_Union[Exists, _Mapping]] = ..., + listStateGet: _Optional[_Union[ListStateGet, _Mapping]] = ..., + listStatePut: _Optional[_Union[ListStatePut, _Mapping]] = ..., + appendValue: _Optional[_Union[AppendValue, _Mapping]] = ..., + appendList: _Optional[_Union[AppendList, _Mapping]] = ..., + clear: _Optional[_Union[Clear, _Mapping]] = ..., ) -> None: ... +class SetImplicitKey(_message.Message): + __slots__ = ("key",) + KEY_FIELD_NUMBER: _ClassVar[int] + key: bytes + def __init__(self, key: _Optional[bytes] = ...) -> None: ... + +class RemoveImplicitKey(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + +class Exists(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + +class Get(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + class ValueStateUpdate(_message.Message): - __slots__ = ["value"] - VALUE_FIELD_NUMBER: ClassVar[int] + __slots__ = ("value",) + VALUE_FIELD_NUMBER: _ClassVar[int] value: bytes - def __init__(self, value: Optional[bytes] = ...) -> None: ... + def __init__(self, value: _Optional[bytes] = ...) -> None: ... -class HandleState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): +class Clear(_message.Message): __slots__ = () + def __init__(self) -> None: ... + +class ListStateGet(_message.Message): + __slots__ = ("iteratorId",) + ITERATORID_FIELD_NUMBER: _ClassVar[int] + iteratorId: str + def __init__(self, iteratorId: _Optional[str] = ...) -> None: ... + +class ListStatePut(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + +class AppendValue(_message.Message): + __slots__ = ("value",) + VALUE_FIELD_NUMBER: _ClassVar[int] + value: bytes + def __init__(self, value: _Optional[bytes] = ...) -> None: ... + +class AppendList(_message.Message): + __slots__ = () + def __init__(self) -> None: ... + +class SetHandleState(_message.Message): + __slots__ = ("state",) + STATE_FIELD_NUMBER: _ClassVar[int] + state: HandleState + def __init__(self, state: _Optional[_Union[HandleState, str]] = ...) -> None: ... + +class TTLConfig(_message.Message): + __slots__ = ("durationMs",) + DURATIONMS_FIELD_NUMBER: _ClassVar[int] + durationMs: int + def __init__(self, durationMs: _Optional[int] = ...) -> None: ... diff --git a/python/pyspark/sql/streaming/list_state_client.py b/python/pyspark/sql/streaming/list_state_client.py new file mode 100644 index 0000000000000..93306eca425eb --- /dev/null +++ b/python/pyspark/sql/streaming/list_state_client.py @@ -0,0 +1,187 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from typing import Dict, Iterator, List, Union, cast, Tuple + +from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient +from pyspark.sql.types import StructType, TYPE_CHECKING, _parse_datatype_string +from pyspark.errors import PySparkRuntimeError +import uuid + +if TYPE_CHECKING: + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike + +__all__ = ["ListStateClient"] + + +class ListStateClient: + def __init__(self, stateful_processor_api_client: StatefulProcessorApiClient) -> None: + self._stateful_processor_api_client = stateful_processor_api_client + # A dictionary to store the mapping between list state name and a tuple of pandas DataFrame + # and the index of the last row that was read. + self.pandas_df_dict: Dict[str, Tuple["PandasDataFrameLike", int]] = {} + + def exists(self, state_name: str) -> bool: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + exists_call = stateMessage.Exists() + list_state_call = stateMessage.ListStateCall(stateName=state_name, exists=exists_call) + state_variable_request = stateMessage.StateVariableRequest(listStateCall=list_state_call) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status == 0: + return True + elif status == 2: + # Expect status code is 2 when state variable doesn't have a value. + return False + else: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError( + f"Error checking value state exists: " f"{response_message[1]}" + ) + + def get(self, state_name: str, iterator_id: str) -> Tuple: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + if iterator_id in self.pandas_df_dict: + # If the state is already in the dictionary, return the next row. + pandas_df, index = self.pandas_df_dict[iterator_id] + else: + # If the state is not in the dictionary, fetch the state from the server. + get_call = stateMessage.ListStateGet(iteratorId=iterator_id) + list_state_call = stateMessage.ListStateCall( + stateName=state_name, listStateGet=get_call + ) + state_variable_request = stateMessage.StateVariableRequest( + listStateCall=list_state_call + ) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status == 0: + iterator = self._stateful_processor_api_client._read_arrow_state() + batch = next(iterator) + pandas_df = batch.to_pandas() + index = 0 + else: + raise StopIteration() + + new_index = index + 1 + if new_index < len(pandas_df): + # Update the index in the dictionary. + self.pandas_df_dict[iterator_id] = (pandas_df, new_index) + else: + # If the index is at the end of the DataFrame, remove the state from the dictionary. + self.pandas_df_dict.pop(iterator_id, None) + pandas_row = pandas_df.iloc[index] + return tuple(pandas_row) + + def append_value(self, state_name: str, schema: Union[StructType, str], value: Tuple) -> None: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + if isinstance(schema, str): + schema = cast(StructType, _parse_datatype_string(schema)) + bytes = self._stateful_processor_api_client._serialize_to_bytes(schema, value) + append_value_call = stateMessage.AppendValue(value=bytes) + list_state_call = stateMessage.ListStateCall( + stateName=state_name, appendValue=append_value_call + ) + state_variable_request = stateMessage.StateVariableRequest(listStateCall=list_state_call) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status != 0: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError(f"Error updating value state: " f"{response_message[1]}") + + def append_list( + self, state_name: str, schema: Union[StructType, str], values: List[Tuple] + ) -> None: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + if isinstance(schema, str): + schema = cast(StructType, _parse_datatype_string(schema)) + append_list_call = stateMessage.AppendList() + list_state_call = stateMessage.ListStateCall( + stateName=state_name, appendList=append_list_call + ) + state_variable_request = stateMessage.StateVariableRequest(listStateCall=list_state_call) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + + self._stateful_processor_api_client._send_arrow_state(schema, values) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status != 0: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError(f"Error updating value state: " f"{response_message[1]}") + + def put(self, state_name: str, schema: Union[StructType, str], values: List[Tuple]) -> None: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + if isinstance(schema, str): + schema = cast(StructType, _parse_datatype_string(schema)) + put_call = stateMessage.ListStatePut() + list_state_call = stateMessage.ListStateCall(stateName=state_name, listStatePut=put_call) + state_variable_request = stateMessage.StateVariableRequest(listStateCall=list_state_call) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + + self._stateful_processor_api_client._send_arrow_state(schema, values) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status != 0: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError(f"Error updating value state: " f"{response_message[1]}") + + def clear(self, state_name: str) -> None: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + clear_call = stateMessage.Clear() + list_state_call = stateMessage.ListStateCall(stateName=state_name, clear=clear_call) + state_variable_request = stateMessage.StateVariableRequest(listStateCall=list_state_call) + message = stateMessage.StateRequest(stateVariableRequest=state_variable_request) + + self._stateful_processor_api_client._send_proto_message(message.SerializeToString()) + response_message = self._stateful_processor_api_client._receive_proto_message() + status = response_message[0] + if status != 0: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError(f"Error clearing value state: " f"{response_message[1]}") + + +class ListStateIterator: + def __init__(self, list_state_client: ListStateClient, state_name: str): + self.list_state_client = list_state_client + self.state_name = state_name + # Generate a unique identifier for the iterator to make sure iterators from the same + # list state do not interfere with each other. + self.iterator_id = str(uuid.uuid4()) + + def __iter__(self) -> Iterator[Tuple]: + return self + + def __next__(self) -> Tuple: + return self.list_state_client.get(self.state_name, self.iterator_id) diff --git a/python/pyspark/sql/streaming/stateful_processor.py b/python/pyspark/sql/streaming/stateful_processor.py index 9045c81e287cd..0011b62132ade 100644 --- a/python/pyspark/sql/streaming/stateful_processor.py +++ b/python/pyspark/sql/streaming/stateful_processor.py @@ -16,12 +16,12 @@ # from abc import ABC, abstractmethod -from typing import Any, TYPE_CHECKING, Iterator, Optional, Union, cast +from typing import Any, List, TYPE_CHECKING, Iterator, Optional, Union, Tuple -from pyspark.sql import Row from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient +from pyspark.sql.streaming.list_state_client import ListStateClient, ListStateIterator from pyspark.sql.streaming.value_state_client import ValueStateClient -from pyspark.sql.types import StructType, _create_row, _parse_datatype_string +from pyspark.sql.types import StructType if TYPE_CHECKING: from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike @@ -50,19 +50,11 @@ def exists(self) -> bool: """ return self._value_state_client.exists(self._state_name) - def get(self) -> Optional[Row]: + def get(self) -> Optional[Tuple]: """ Get the state value if it exists. Returns None if the state variable does not have a value. """ - value = self._value_state_client.get(self._state_name) - if value is None: - return None - schema = self.schema - if isinstance(schema, str): - schema = cast(StructType, _parse_datatype_string(schema)) - # Create the Row using the values and schema fields - row = _create_row(schema.fieldNames(), value) - return row + return self._value_state_client.get(self._state_name) def update(self, new_value: Any) -> None: """ @@ -77,6 +69,58 @@ def clear(self) -> None: self._value_state_client.clear(self._state_name) +class ListState: + """ + Class used for arbitrary stateful operations with transformWithState to capture list value + state. + + .. versionadded:: 4.0.0 + """ + + def __init__( + self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str] + ) -> None: + self._list_state_client = list_state_client + self._state_name = state_name + self.schema = schema + + def exists(self) -> bool: + """ + Whether list state exists or not. + """ + return self._list_state_client.exists(self._state_name) + + def get(self) -> Iterator[Tuple]: + """ + Get list state with an iterator. + """ + return ListStateIterator(self._list_state_client, self._state_name) + + def put(self, new_state: List[Tuple]) -> None: + """ + Update the values of the list state. + """ + self._list_state_client.put(self._state_name, self.schema, new_state) + + def append_value(self, new_state: Tuple) -> None: + """ + Append a new value to the list state. + """ + self._list_state_client.append_value(self._state_name, self.schema, new_state) + + def append_list(self, new_state: List[Tuple]) -> None: + """ + Append a list of new values to the list state. + """ + self._list_state_client.append_list(self._state_name, self.schema, new_state) + + def clear(self) -> None: + """ + Remove this state. + """ + self._list_state_client.clear(self._state_name) + + class StatefulProcessorHandle: """ Represents the operation handle provided to the stateful processor used in transformWithState @@ -112,6 +156,23 @@ def getValueState( self.stateful_processor_api_client.get_value_state(state_name, schema, ttl_duration_ms) return ValueState(ValueStateClient(self.stateful_processor_api_client), state_name, schema) + def getListState(self, state_name: str, schema: Union[StructType, str]) -> ListState: + """ + Function to create new or return existing single value state variable of given type. + The user must ensure to call this function only within the `init()` method of the + :class:`StatefulProcessor`. + + Parameters + ---------- + state_name : str + name of the state variable + schema : :class:`pyspark.sql.types.DataType` or str + The schema of the state variable. The value can be either a + :class:`pyspark.sql.types.DataType` object or a DDL-formatted type string. + """ + self.stateful_processor_api_client.get_list_state(state_name, schema) + return ListState(ListStateClient(self.stateful_processor_api_client), state_name, schema) + class StatefulProcessor(ABC): """ diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py b/python/pyspark/sql/streaming/stateful_processor_api_client.py index 9703aa17d3474..2a5e55159e766 100644 --- a/python/pyspark/sql/streaming/stateful_processor_api_client.py +++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py @@ -17,10 +17,16 @@ from enum import Enum import os import socket -from typing import Any, Union, Optional, cast, Tuple +from typing import Any, List, Union, Optional, cast, Tuple from pyspark.serializers import write_int, read_int, UTF8Deserializer -from pyspark.sql.types import StructType, _parse_datatype_string, Row +from pyspark.sql.pandas.serializers import ArrowStreamSerializer +from pyspark.sql.types import ( + StructType, + _parse_datatype_string, + Row, +) +from pyspark.sql.pandas.types import convert_pandas_using_numpy_type from pyspark.sql.utils import has_numpy from pyspark.serializers import CPickleSerializer from pyspark.errors import PySparkRuntimeError @@ -46,6 +52,7 @@ def __init__(self, state_server_port: int, key_schema: StructType) -> None: self.handle_state = StatefulProcessorHandleState.CREATED self.utf8_deserializer = UTF8Deserializer() self.pickleSer = CPickleSerializer() + self.serializer = ArrowStreamSerializer() def set_handle_state(self, state: StatefulProcessorHandleState) -> None: import pyspark.sql.streaming.StateMessage_pb2 as stateMessage @@ -124,6 +131,25 @@ def get_value_state( # TODO(SPARK-49233): Classify user facing errors. raise PySparkRuntimeError(f"Error initializing value state: " f"{response_message[1]}") + def get_list_state(self, state_name: str, schema: Union[StructType, str]) -> None: + import pyspark.sql.streaming.StateMessage_pb2 as stateMessage + + if isinstance(schema, str): + schema = cast(StructType, _parse_datatype_string(schema)) + + state_call_command = stateMessage.StateCallCommand() + state_call_command.stateName = state_name + state_call_command.schema = schema.json() + call = stateMessage.StatefulProcessorCall(getListState=state_call_command) + message = stateMessage.StateRequest(statefulProcessorCall=call) + + self._send_proto_message(message.SerializeToString()) + response_message = self._receive_proto_message() + status = response_message[0] + if status != 0: + # TODO(SPARK-49233): Classify user facing errors. + raise PySparkRuntimeError(f"Error initializing value state: " f"{response_message[1]}") + def _send_proto_message(self, message: bytes) -> None: # Writing zero here to indicate message version. This allows us to evolve the message # format or even changing the message protocol in the future. @@ -168,3 +194,18 @@ def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes: def _deserialize_from_bytes(self, value: bytes) -> Any: return self.pickleSer.loads(value) + + def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None: + import pyarrow as pa + import pandas as pd + + column_names = [field.name for field in schema.fields] + pandas_df = convert_pandas_using_numpy_type( + pd.DataFrame(state, columns=column_names), schema + ) + batch = pa.RecordBatch.from_pandas(pandas_df) + self.serializer.dump_stream(iter([batch]), self.sockfile) + self.sockfile.flush() + + def _read_arrow_state(self) -> Any: + return self.serializer.load_stream(self.sockfile) diff --git a/python/pyspark/sql/streaming/value_state_client.py b/python/pyspark/sql/streaming/value_state_client.py index e902f70cb40a5..3fe32bcc5235c 100644 --- a/python/pyspark/sql/streaming/value_state_client.py +++ b/python/pyspark/sql/streaming/value_state_client.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Union, cast, Tuple +from typing import Union, cast, Tuple, Optional from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient from pyspark.sql.types import StructType, _parse_datatype_string @@ -49,7 +49,7 @@ def exists(self, state_name: str) -> bool: f"Error checking value state exists: " f"{response_message[1]}" ) - def get(self, state_name: str) -> Any: + def get(self, state_name: str) -> Optional[Tuple]: import pyspark.sql.streaming.StateMessage_pb2 as stateMessage get_call = stateMessage.Get() @@ -63,8 +63,8 @@ def get(self, state_name: str) -> Any: if status == 0: if len(response_message[2]) == 0: return None - row = self._stateful_processor_api_client._deserialize_from_bytes(response_message[2]) - return row + data = self._stateful_processor_api_client._deserialize_from_bytes(response_message[2]) + return tuple(data) else: # TODO(SPARK-49233): Classify user facing errors. raise PySparkRuntimeError(f"Error getting value state: " f"{response_message[1]}") diff --git a/python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py index 8ad24704de3a4..99333ae6f5c26 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py @@ -59,6 +59,7 @@ def conf(cls): "spark.sql.streaming.stateStore.providerClass", "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider", ) + cfg.set("spark.sql.execution.arrow.transformWithStateInPandas.maxRecordsPerBatch", "2") return cfg def _prepare_input_data(self, input_path, col1, col2): @@ -211,6 +212,15 @@ def test_transform_with_state_in_pandas_query_restarts(self): Row(id="1", countAsString="2"), } + def test_transform_with_state_in_pandas_list_state(self): + def check_results(batch_df, _): + assert set(batch_df.sort("id").collect()) == { + Row(id="0", countAsString="2"), + Row(id="1", countAsString="2"), + } + + self._test_transform_with_state_in_pandas_basic(ListStateProcessor(), check_results, True) + # test value state with ttl has the same behavior as value state when # state doesn't expire. def test_value_state_ttl_basic(self): @@ -394,6 +404,59 @@ def close(self) -> None: pass +class ListStateProcessor(StatefulProcessor): + # Dict to store the expected results. The key represents the grouping key string, and the value + # is a dictionary of pandas dataframe index -> expected temperature value. Since we set + # maxRecordsPerBatch to 2, we expect the pandas dataframe dictionary to have 2 entries. + dict = {0: 120, 1: 20} + + def init(self, handle: StatefulProcessorHandle) -> None: + state_schema = StructType([StructField("temperature", IntegerType(), True)]) + self.list_state1 = handle.getListState("listState1", state_schema) + self.list_state2 = handle.getListState("listState2", state_schema) + + def handleInputRows(self, key, rows) -> Iterator[pd.DataFrame]: + count = 0 + for pdf in rows: + list_state_rows = [(120,), (20,)] + self.list_state1.put(list_state_rows) + self.list_state2.put(list_state_rows) + self.list_state1.append_value((111,)) + self.list_state2.append_value((222,)) + self.list_state1.append_list(list_state_rows) + self.list_state2.append_list(list_state_rows) + pdf_count = pdf.count() + count += pdf_count.get("temperature") + iter1 = self.list_state1.get() + iter2 = self.list_state2.get() + # Mixing the iterator to test it we can resume from the correct point + assert next(iter1)[0] == self.dict[0] + assert next(iter2)[0] == self.dict[0] + assert next(iter1)[0] == self.dict[1] + assert next(iter2)[0] == self.dict[1] + # Get another iterator for list_state1 to test if the 2 iterators (iter1 and iter3) don't + # interfere with each other. + iter3 = self.list_state1.get() + assert next(iter3)[0] == self.dict[0] + assert next(iter3)[0] == self.dict[1] + # the second arrow batch should contain the appended value 111 for list_state1 and + # 222 for list_state2 + assert next(iter1)[0] == 111 + assert next(iter2)[0] == 222 + assert next(iter3)[0] == 111 + # since we put another 2 rows after 111/222, check them here + assert next(iter1)[0] == self.dict[0] + assert next(iter2)[0] == self.dict[0] + assert next(iter3)[0] == self.dict[0] + assert next(iter1)[0] == self.dict[1] + assert next(iter2)[0] == self.dict[1] + assert next(iter3)[0] == self.dict[1] + yield pd.DataFrame({"id": key, "countAsString": str(count)}) + + def close(self) -> None: + pass + + class TransformWithStateInPandasTests(TransformWithStateInPandasTestsMixin, ReusedSQLTestCase): pass diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 9d51afd064d10..c9c227a21cfff 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3217,6 +3217,14 @@ object SQLConf { .intConf .createWithDefault(10000) + val ARROW_TRANSFORM_WITH_STATE_IN_PANDAS_MAX_RECORDS_PER_BATCH = + buildConf("spark.sql.execution.arrow.transformWithStateInPandas.maxRecordsPerBatch") + .doc("When using TransformWithStateInPandas, limit the maximum number of state records " + + "that can be written to a single ArrowRecordBatch in memory.") + .version("4.0.0") + .intConf + .createWithDefault(10000) + val ARROW_EXECUTION_USE_LARGE_VAR_TYPES = buildConf("spark.sql.execution.arrow.useLargeVarTypes") .doc("When using Apache Arrow, use large variable width vectors for string and binary " + @@ -5899,6 +5907,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def arrowMaxRecordsPerBatch: Int = getConf(ARROW_EXECUTION_MAX_RECORDS_PER_BATCH) + def arrowTransformWithStateInPandasMaxRecordsPerBatch: Int = + getConf(ARROW_TRANSFORM_WITH_STATE_IN_PANDAS_MAX_RECORDS_PER_BATCH) + def arrowUseLargeVarTypes: Boolean = getConf(ARROW_EXECUTION_USE_LARGE_VAR_TYPES) def pandasUDFBufferSize: Int = getConf(PANDAS_UDF_BUFFER_SIZE) diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/StateMessage.proto b/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/StateMessage.proto index 1ff90f27e173a..63728216ded1e 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/StateMessage.proto +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/StateMessage.proto @@ -46,6 +46,7 @@ message StatefulProcessorCall { message StateVariableRequest { oneof method { ValueStateCall valueStateCall = 1; + ListStateCall listStateCall = 2; } } @@ -72,6 +73,18 @@ message ValueStateCall { } } +message ListStateCall { + string stateName = 1; + oneof method { + Exists exists = 2; + ListStateGet listStateGet = 3; + ListStatePut listStatePut = 4; + AppendValue appendValue = 5; + AppendList appendList = 6; + Clear clear = 7; + } +} + message SetImplicitKey { bytes key = 1; } @@ -92,6 +105,20 @@ message ValueStateUpdate { message Clear { } +message ListStateGet { + string iteratorId = 1; +} + +message ListStatePut { +} + +message AppendValue { + bytes value = 1; +} + +message AppendList { +} + enum HandleState { CREATED = 0; INITIALIZED = 1; diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/state/StateMessage.java b/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/state/StateMessage.java index 4fbb20be05b7b..d6d56dd732775 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/state/StateMessage.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/streaming/state/StateMessage.java @@ -3462,6 +3462,21 @@ public interface StateVariableRequestOrBuilder extends */ org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCallOrBuilder getValueStateCallOrBuilder(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return Whether the listStateCall field is set. + */ + boolean hasListStateCall(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return The listStateCall. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getListStateCall(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder getListStateCallOrBuilder(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.StateVariableRequest.MethodCase getMethodCase(); } /** @@ -3510,6 +3525,7 @@ public enum MethodCase implements com.google.protobuf.Internal.EnumLite, com.google.protobuf.AbstractMessage.InternalOneOfEnum { VALUESTATECALL(1), + LISTSTATECALL(2), METHOD_NOT_SET(0); private final int value; private MethodCase(int value) { @@ -3528,6 +3544,7 @@ public static MethodCase valueOf(int value) { public static MethodCase forNumber(int value) { switch (value) { case 1: return VALUESTATECALL; + case 2: return LISTSTATECALL; case 0: return METHOD_NOT_SET; default: return null; } @@ -3574,6 +3591,37 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCal return org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCall.getDefaultInstance(); } + public static final int LISTSTATECALL_FIELD_NUMBER = 2; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return Whether the listStateCall field is set. + */ + @java.lang.Override + public boolean hasListStateCall() { + return methodCase_ == 2; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return The listStateCall. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getListStateCall() { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder getListStateCallOrBuilder() { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + private byte memoizedIsInitialized = -1; @java.lang.Override public final boolean isInitialized() { @@ -3591,6 +3639,9 @@ public void writeTo(com.google.protobuf.CodedOutputStream output) if (methodCase_ == 1) { output.writeMessage(1, (org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCall) method_); } + if (methodCase_ == 2) { + output.writeMessage(2, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_); + } getUnknownFields().writeTo(output); } @@ -3604,6 +3655,10 @@ public int getSerializedSize() { size += com.google.protobuf.CodedOutputStream .computeMessageSize(1, (org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCall) method_); } + if (methodCase_ == 2) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(2, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_); + } size += getUnknownFields().getSerializedSize(); memoizedSize = size; return size; @@ -3625,6 +3680,10 @@ public boolean equals(final java.lang.Object obj) { if (!getValueStateCall() .equals(other.getValueStateCall())) return false; break; + case 2: + if (!getListStateCall() + .equals(other.getListStateCall())) return false; + break; case 0: default: } @@ -3644,6 +3703,10 @@ public int hashCode() { hash = (37 * hash) + VALUESTATECALL_FIELD_NUMBER; hash = (53 * hash) + getValueStateCall().hashCode(); break; + case 2: + hash = (37 * hash) + LISTSTATECALL_FIELD_NUMBER; + hash = (53 * hash) + getListStateCall().hashCode(); + break; case 0: default: } @@ -3778,6 +3841,9 @@ public Builder clear() { if (valueStateCallBuilder_ != null) { valueStateCallBuilder_.clear(); } + if (listStateCallBuilder_ != null) { + listStateCallBuilder_.clear(); + } methodCase_ = 0; method_ = null; return this; @@ -3813,6 +3879,13 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.StateVariable result.method_ = valueStateCallBuilder_.build(); } } + if (methodCase_ == 2) { + if (listStateCallBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = listStateCallBuilder_.build(); + } + } result.methodCase_ = methodCase_; onBuilt(); return result; @@ -3867,6 +3940,10 @@ public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMes mergeValueStateCall(other.getValueStateCall()); break; } + case LISTSTATECALL: { + mergeListStateCall(other.getListStateCall()); + break; + } case METHOD_NOT_SET: { break; } @@ -3904,6 +3981,13 @@ public Builder mergeFrom( methodCase_ = 1; break; } // case 10 + case 18: { + input.readMessage( + getListStateCallFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 2; + break; + } // case 18 default: { if (!super.parseUnknownField(input, extensionRegistry, tag)) { done = true; // was an endgroup tag @@ -4076,6 +4160,148 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCal onChanged();; return valueStateCallBuilder_; } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder> listStateCallBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return Whether the listStateCall field is set. + */ + @java.lang.Override + public boolean hasListStateCall() { + return methodCase_ == 2; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + * @return The listStateCall. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getListStateCall() { + if (listStateCallBuilder_ == null) { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } else { + if (methodCase_ == 2) { + return listStateCallBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + public Builder setListStateCall(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall value) { + if (listStateCallBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + listStateCallBuilder_.setMessage(value); + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + public Builder setListStateCall( + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder builderForValue) { + if (listStateCallBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + listStateCallBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + public Builder mergeListStateCall(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall value) { + if (listStateCallBuilder_ == null) { + if (methodCase_ == 2 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 2) { + listStateCallBuilder_.mergeFrom(value); + } else { + listStateCallBuilder_.setMessage(value); + } + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + public Builder clearListStateCall() { + if (listStateCallBuilder_ == null) { + if (methodCase_ == 2) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 2) { + methodCase_ = 0; + method_ = null; + } + listStateCallBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder getListStateCallBuilder() { + return getListStateCallFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder getListStateCallOrBuilder() { + if ((methodCase_ == 2) && (listStateCallBuilder_ != null)) { + return listStateCallBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateCall listStateCall = 2; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder> + getListStateCallFieldBuilder() { + if (listStateCallBuilder_ == null) { + if (!(methodCase_ == 2)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + listStateCallBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 2; + onChanged();; + return listStateCallBuilder_; + } @java.lang.Override public final Builder setUnknownFields( final com.google.protobuf.UnknownFieldSet unknownFields) { @@ -7482,37 +7708,135 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateCal } - public interface SetImplicitKeyOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + public interface ListStateCallOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.ListStateCall) com.google.protobuf.MessageOrBuilder { /** - * bytes key = 1; - * @return The key. + * string stateName = 1; + * @return The stateName. */ - com.google.protobuf.ByteString getKey(); + java.lang.String getStateName(); + /** + * string stateName = 1; + * @return The bytes for stateName. + */ + com.google.protobuf.ByteString + getStateNameBytes(); + + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return Whether the exists field is set. + */ + boolean hasExists(); + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return The exists. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getExists(); + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder getExistsOrBuilder(); + + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return Whether the listStateGet field is set. + */ + boolean hasListStateGet(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return The listStateGet. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getListStateGet(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder getListStateGetOrBuilder(); + + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return Whether the listStatePut field is set. + */ + boolean hasListStatePut(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return The listStatePut. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getListStatePut(); + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder getListStatePutOrBuilder(); + + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return Whether the appendValue field is set. + */ + boolean hasAppendValue(); + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return The appendValue. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getAppendValue(); + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder getAppendValueOrBuilder(); + + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return Whether the appendList field is set. + */ + boolean hasAppendList(); + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return The appendList. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getAppendList(); + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder getAppendListOrBuilder(); + + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return Whether the clear field is set. + */ + boolean hasClear(); + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return The clear. + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getClear(); + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder getClearOrBuilder(); + + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.MethodCase getMethodCase(); } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.SetImplicitKey} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStateCall} */ - public static final class SetImplicitKey extends + public static final class ListStateCall extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) - SetImplicitKeyOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.ListStateCall) + ListStateCallOrBuilder { private static final long serialVersionUID = 0L; - // Use SetImplicitKey.newBuilder() to construct. - private SetImplicitKey(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use ListStateCall.newBuilder() to construct. + private ListStateCall(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private SetImplicitKey() { - key_ = com.google.protobuf.ByteString.EMPTY; + private ListStateCall() { + stateName_ = ""; } @java.lang.Override @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new SetImplicitKey(); + return new ListStateCall(); } @java.lang.Override @@ -7522,31 +7846,3583 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder.class); } - public static final int KEY_FIELD_NUMBER = 1; - private com.google.protobuf.ByteString key_; - /** - * bytes key = 1; - * @return The key. - */ - @java.lang.Override - public com.google.protobuf.ByteString getKey() { - return key_; - } + private int methodCase_ = 0; + private java.lang.Object method_; + public enum MethodCase + implements com.google.protobuf.Internal.EnumLite, + com.google.protobuf.AbstractMessage.InternalOneOfEnum { + EXISTS(2), + LISTSTATEGET(3), + LISTSTATEPUT(4), + APPENDVALUE(5), + APPENDLIST(6), + CLEAR(7), + METHOD_NOT_SET(0); + private final int value; + private MethodCase(int value) { + this.value = value; + } + /** + * @param value The number of the enum to look for. + * @return The enum associated with the given number. + * @deprecated Use {@link #forNumber(int)} instead. + */ + @java.lang.Deprecated + public static MethodCase valueOf(int value) { + return forNumber(value); + } - private byte memoizedIsInitialized = -1; - @java.lang.Override - public final boolean isInitialized() { + public static MethodCase forNumber(int value) { + switch (value) { + case 2: return EXISTS; + case 3: return LISTSTATEGET; + case 4: return LISTSTATEPUT; + case 5: return APPENDVALUE; + case 6: return APPENDLIST; + case 7: return CLEAR; + case 0: return METHOD_NOT_SET; + default: return null; + } + } + public int getNumber() { + return this.value; + } + }; + + public MethodCase + getMethodCase() { + return MethodCase.forNumber( + methodCase_); + } + + public static final int STATENAME_FIELD_NUMBER = 1; + private volatile java.lang.Object stateName_; + /** + * string stateName = 1; + * @return The stateName. + */ + @java.lang.Override + public java.lang.String getStateName() { + java.lang.Object ref = stateName_; + if (ref instanceof java.lang.String) { + return (java.lang.String) ref; + } else { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + stateName_ = s; + return s; + } + } + /** + * string stateName = 1; + * @return The bytes for stateName. + */ + @java.lang.Override + public com.google.protobuf.ByteString + getStateNameBytes() { + java.lang.Object ref = stateName_; + if (ref instanceof java.lang.String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + stateName_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + + public static final int EXISTS_FIELD_NUMBER = 2; + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return Whether the exists field is set. + */ + @java.lang.Override + public boolean hasExists() { + return methodCase_ == 2; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return The exists. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getExists() { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder getExistsOrBuilder() { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + + public static final int LISTSTATEGET_FIELD_NUMBER = 3; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return Whether the listStateGet field is set. + */ + @java.lang.Override + public boolean hasListStateGet() { + return methodCase_ == 3; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return The listStateGet. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getListStateGet() { + if (methodCase_ == 3) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder getListStateGetOrBuilder() { + if (methodCase_ == 3) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } + + public static final int LISTSTATEPUT_FIELD_NUMBER = 4; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return Whether the listStatePut field is set. + */ + @java.lang.Override + public boolean hasListStatePut() { + return methodCase_ == 4; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return The listStatePut. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getListStatePut() { + if (methodCase_ == 4) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder getListStatePutOrBuilder() { + if (methodCase_ == 4) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } + + public static final int APPENDVALUE_FIELD_NUMBER = 5; + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return Whether the appendValue field is set. + */ + @java.lang.Override + public boolean hasAppendValue() { + return methodCase_ == 5; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return The appendValue. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getAppendValue() { + if (methodCase_ == 5) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder getAppendValueOrBuilder() { + if (methodCase_ == 5) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } + + public static final int APPENDLIST_FIELD_NUMBER = 6; + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return Whether the appendList field is set. + */ + @java.lang.Override + public boolean hasAppendList() { + return methodCase_ == 6; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return The appendList. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getAppendList() { + if (methodCase_ == 6) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder getAppendListOrBuilder() { + if (methodCase_ == 6) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } + + public static final int CLEAR_FIELD_NUMBER = 7; + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return Whether the clear field is set. + */ + @java.lang.Override + public boolean hasClear() { + return methodCase_ == 7; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return The clear. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getClear() { + if (methodCase_ == 7) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder getClearOrBuilder() { + if (methodCase_ == 7) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(stateName_)) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 1, stateName_); + } + if (methodCase_ == 2) { + output.writeMessage(2, (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_); + } + if (methodCase_ == 3) { + output.writeMessage(3, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_); + } + if (methodCase_ == 4) { + output.writeMessage(4, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_); + } + if (methodCase_ == 5) { + output.writeMessage(5, (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_); + } + if (methodCase_ == 6) { + output.writeMessage(6, (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_); + } + if (methodCase_ == 7) { + output.writeMessage(7, (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_); + } + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(stateName_)) { + size += com.google.protobuf.GeneratedMessageV3.computeStringSize(1, stateName_); + } + if (methodCase_ == 2) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(2, (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_); + } + if (methodCase_ == 3) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(3, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_); + } + if (methodCase_ == 4) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(4, (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_); + } + if (methodCase_ == 5) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(5, (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_); + } + if (methodCase_ == 6) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(6, (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_); + } + if (methodCase_ == 7) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(7, (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_); + } + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall)) { + return super.equals(obj); + } + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall other = (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) obj; + + if (!getStateName() + .equals(other.getStateName())) return false; + if (!getMethodCase().equals(other.getMethodCase())) return false; + switch (methodCase_) { + case 2: + if (!getExists() + .equals(other.getExists())) return false; + break; + case 3: + if (!getListStateGet() + .equals(other.getListStateGet())) return false; + break; + case 4: + if (!getListStatePut() + .equals(other.getListStatePut())) return false; + break; + case 5: + if (!getAppendValue() + .equals(other.getAppendValue())) return false; + break; + case 6: + if (!getAppendList() + .equals(other.getAppendList())) return false; + break; + case 7: + if (!getClear() + .equals(other.getClear())) return false; + break; + case 0: + default: + } + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (37 * hash) + STATENAME_FIELD_NUMBER; + hash = (53 * hash) + getStateName().hashCode(); + switch (methodCase_) { + case 2: + hash = (37 * hash) + EXISTS_FIELD_NUMBER; + hash = (53 * hash) + getExists().hashCode(); + break; + case 3: + hash = (37 * hash) + LISTSTATEGET_FIELD_NUMBER; + hash = (53 * hash) + getListStateGet().hashCode(); + break; + case 4: + hash = (37 * hash) + LISTSTATEPUT_FIELD_NUMBER; + hash = (53 * hash) + getListStatePut().hashCode(); + break; + case 5: + hash = (37 * hash) + APPENDVALUE_FIELD_NUMBER; + hash = (53 * hash) + getAppendValue().hashCode(); + break; + case 6: + hash = (37 * hash) + APPENDLIST_FIELD_NUMBER; + hash = (53 * hash) + getAppendList().hashCode(); + break; + case 7: + hash = (37 * hash) + CLEAR_FIELD_NUMBER; + hash = (53 * hash) + getClear().hashCode(); + break; + case 0: + default: + } + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStateCall} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.ListStateCall) + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCallOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.Builder.class); + } + + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + stateName_ = ""; + + if (existsBuilder_ != null) { + existsBuilder_.clear(); + } + if (listStateGetBuilder_ != null) { + listStateGetBuilder_.clear(); + } + if (listStatePutBuilder_ != null) { + listStatePutBuilder_.clear(); + } + if (appendValueBuilder_ != null) { + appendValueBuilder_.clear(); + } + if (appendListBuilder_ != null) { + appendListBuilder_.clear(); + } + if (clearBuilder_ != null) { + clearBuilder_.clear(); + } + methodCase_ = 0; + method_ = null; + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance(); + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall result = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall(this); + result.stateName_ = stateName_; + if (methodCase_ == 2) { + if (existsBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = existsBuilder_.build(); + } + } + if (methodCase_ == 3) { + if (listStateGetBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = listStateGetBuilder_.build(); + } + } + if (methodCase_ == 4) { + if (listStatePutBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = listStatePutBuilder_.build(); + } + } + if (methodCase_ == 5) { + if (appendValueBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = appendValueBuilder_.build(); + } + } + if (methodCase_ == 6) { + if (appendListBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = appendListBuilder_.build(); + } + } + if (methodCase_ == 7) { + if (clearBuilder_ == null) { + result.method_ = method_; + } else { + result.method_ = clearBuilder_.build(); + } + } + result.methodCase_ = methodCase_; + onBuilt(); + return result; + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall.getDefaultInstance()) return this; + if (!other.getStateName().isEmpty()) { + stateName_ = other.stateName_; + onChanged(); + } + switch (other.getMethodCase()) { + case EXISTS: { + mergeExists(other.getExists()); + break; + } + case LISTSTATEGET: { + mergeListStateGet(other.getListStateGet()); + break; + } + case LISTSTATEPUT: { + mergeListStatePut(other.getListStatePut()); + break; + } + case APPENDVALUE: { + mergeAppendValue(other.getAppendValue()); + break; + } + case APPENDLIST: { + mergeAppendList(other.getAppendList()); + break; + } + case CLEAR: { + mergeClear(other.getClear()); + break; + } + case METHOD_NOT_SET: { + break; + } + } + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + case 10: { + stateName_ = input.readStringRequireUtf8(); + + break; + } // case 10 + case 18: { + input.readMessage( + getExistsFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 2; + break; + } // case 18 + case 26: { + input.readMessage( + getListStateGetFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 3; + break; + } // case 26 + case 34: { + input.readMessage( + getListStatePutFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 4; + break; + } // case 34 + case 42: { + input.readMessage( + getAppendValueFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 5; + break; + } // case 42 + case 50: { + input.readMessage( + getAppendListFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 6; + break; + } // case 50 + case 58: { + input.readMessage( + getClearFieldBuilder().getBuilder(), + extensionRegistry); + methodCase_ = 7; + break; + } // case 58 + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + private int methodCase_ = 0; + private java.lang.Object method_; + public MethodCase + getMethodCase() { + return MethodCase.forNumber( + methodCase_); + } + + public Builder clearMethod() { + methodCase_ = 0; + method_ = null; + onChanged(); + return this; + } + + + private java.lang.Object stateName_ = ""; + /** + * string stateName = 1; + * @return The stateName. + */ + public java.lang.String getStateName() { + java.lang.Object ref = stateName_; + if (!(ref instanceof java.lang.String)) { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + stateName_ = s; + return s; + } else { + return (java.lang.String) ref; + } + } + /** + * string stateName = 1; + * @return The bytes for stateName. + */ + public com.google.protobuf.ByteString + getStateNameBytes() { + java.lang.Object ref = stateName_; + if (ref instanceof String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + stateName_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + /** + * string stateName = 1; + * @param value The stateName to set. + * @return This builder for chaining. + */ + public Builder setStateName( + java.lang.String value) { + if (value == null) { + throw new NullPointerException(); + } + + stateName_ = value; + onChanged(); + return this; + } + /** + * string stateName = 1; + * @return This builder for chaining. + */ + public Builder clearStateName() { + + stateName_ = getDefaultInstance().getStateName(); + onChanged(); + return this; + } + /** + * string stateName = 1; + * @param value The bytes for stateName to set. + * @return This builder for chaining. + */ + public Builder setStateNameBytes( + com.google.protobuf.ByteString value) { + if (value == null) { + throw new NullPointerException(); + } + checkByteStringIsUtf8(value); + + stateName_ = value; + onChanged(); + return this; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder> existsBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return Whether the exists field is set. + */ + @java.lang.Override + public boolean hasExists() { + return methodCase_ == 2; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + * @return The exists. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getExists() { + if (existsBuilder_ == null) { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } else { + if (methodCase_ == 2) { + return existsBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + public Builder setExists(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists value) { + if (existsBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + existsBuilder_.setMessage(value); + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + public Builder setExists( + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder builderForValue) { + if (existsBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + existsBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + public Builder mergeExists(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists value) { + if (existsBuilder_ == null) { + if (methodCase_ == 2 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 2) { + existsBuilder_.mergeFrom(value); + } else { + existsBuilder_.setMessage(value); + } + } + methodCase_ = 2; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + public Builder clearExists() { + if (existsBuilder_ == null) { + if (methodCase_ == 2) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 2) { + methodCase_ = 0; + method_ = null; + } + existsBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder getExistsBuilder() { + return getExistsFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder getExistsOrBuilder() { + if ((methodCase_ == 2) && (existsBuilder_ != null)) { + return existsBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 2) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.Exists exists = 2; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder> + getExistsFieldBuilder() { + if (existsBuilder_ == null) { + if (!(methodCase_ == 2)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + existsBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 2; + onChanged();; + return existsBuilder_; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder> listStateGetBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return Whether the listStateGet field is set. + */ + @java.lang.Override + public boolean hasListStateGet() { + return methodCase_ == 3; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + * @return The listStateGet. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getListStateGet() { + if (listStateGetBuilder_ == null) { + if (methodCase_ == 3) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } else { + if (methodCase_ == 3) { + return listStateGetBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + public Builder setListStateGet(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet value) { + if (listStateGetBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + listStateGetBuilder_.setMessage(value); + } + methodCase_ = 3; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + public Builder setListStateGet( + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder builderForValue) { + if (listStateGetBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + listStateGetBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 3; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + public Builder mergeListStateGet(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet value) { + if (listStateGetBuilder_ == null) { + if (methodCase_ == 3 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 3) { + listStateGetBuilder_.mergeFrom(value); + } else { + listStateGetBuilder_.setMessage(value); + } + } + methodCase_ = 3; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + public Builder clearListStateGet() { + if (listStateGetBuilder_ == null) { + if (methodCase_ == 3) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 3) { + methodCase_ = 0; + method_ = null; + } + listStateGetBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder getListStateGetBuilder() { + return getListStateGetFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder getListStateGetOrBuilder() { + if ((methodCase_ == 3) && (listStateGetBuilder_ != null)) { + return listStateGetBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 3) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStateGet listStateGet = 3; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder> + getListStateGetFieldBuilder() { + if (listStateGetBuilder_ == null) { + if (!(methodCase_ == 3)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); + } + listStateGetBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 3; + onChanged();; + return listStateGetBuilder_; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder> listStatePutBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return Whether the listStatePut field is set. + */ + @java.lang.Override + public boolean hasListStatePut() { + return methodCase_ == 4; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + * @return The listStatePut. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getListStatePut() { + if (listStatePutBuilder_ == null) { + if (methodCase_ == 4) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } else { + if (methodCase_ == 4) { + return listStatePutBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + public Builder setListStatePut(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut value) { + if (listStatePutBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + listStatePutBuilder_.setMessage(value); + } + methodCase_ = 4; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + public Builder setListStatePut( + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder builderForValue) { + if (listStatePutBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + listStatePutBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 4; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + public Builder mergeListStatePut(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut value) { + if (listStatePutBuilder_ == null) { + if (methodCase_ == 4 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 4) { + listStatePutBuilder_.mergeFrom(value); + } else { + listStatePutBuilder_.setMessage(value); + } + } + methodCase_ = 4; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + public Builder clearListStatePut() { + if (listStatePutBuilder_ == null) { + if (methodCase_ == 4) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 4) { + methodCase_ = 0; + method_ = null; + } + listStatePutBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder getListStatePutBuilder() { + return getListStatePutFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder getListStatePutOrBuilder() { + if ((methodCase_ == 4) && (listStatePutBuilder_ != null)) { + return listStatePutBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 4) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.ListStatePut listStatePut = 4; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder> + getListStatePutFieldBuilder() { + if (listStatePutBuilder_ == null) { + if (!(methodCase_ == 4)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); + } + listStatePutBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 4; + onChanged();; + return listStatePutBuilder_; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder> appendValueBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return Whether the appendValue field is set. + */ + @java.lang.Override + public boolean hasAppendValue() { + return methodCase_ == 5; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + * @return The appendValue. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getAppendValue() { + if (appendValueBuilder_ == null) { + if (methodCase_ == 5) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } else { + if (methodCase_ == 5) { + return appendValueBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + public Builder setAppendValue(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue value) { + if (appendValueBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + appendValueBuilder_.setMessage(value); + } + methodCase_ = 5; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + public Builder setAppendValue( + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder builderForValue) { + if (appendValueBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + appendValueBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 5; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + public Builder mergeAppendValue(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue value) { + if (appendValueBuilder_ == null) { + if (methodCase_ == 5 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 5) { + appendValueBuilder_.mergeFrom(value); + } else { + appendValueBuilder_.setMessage(value); + } + } + methodCase_ = 5; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + public Builder clearAppendValue() { + if (appendValueBuilder_ == null) { + if (methodCase_ == 5) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 5) { + methodCase_ = 0; + method_ = null; + } + appendValueBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder getAppendValueBuilder() { + return getAppendValueFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder getAppendValueOrBuilder() { + if ((methodCase_ == 5) && (appendValueBuilder_ != null)) { + return appendValueBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 5) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendValue appendValue = 5; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder> + getAppendValueFieldBuilder() { + if (appendValueBuilder_ == null) { + if (!(methodCase_ == 5)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); + } + appendValueBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 5; + onChanged();; + return appendValueBuilder_; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder> appendListBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return Whether the appendList field is set. + */ + @java.lang.Override + public boolean hasAppendList() { + return methodCase_ == 6; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + * @return The appendList. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getAppendList() { + if (appendListBuilder_ == null) { + if (methodCase_ == 6) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } else { + if (methodCase_ == 6) { + return appendListBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + public Builder setAppendList(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList value) { + if (appendListBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + appendListBuilder_.setMessage(value); + } + methodCase_ = 6; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + public Builder setAppendList( + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder builderForValue) { + if (appendListBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + appendListBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 6; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + public Builder mergeAppendList(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList value) { + if (appendListBuilder_ == null) { + if (methodCase_ == 6 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 6) { + appendListBuilder_.mergeFrom(value); + } else { + appendListBuilder_.setMessage(value); + } + } + methodCase_ = 6; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + public Builder clearAppendList() { + if (appendListBuilder_ == null) { + if (methodCase_ == 6) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 6) { + methodCase_ = 0; + method_ = null; + } + appendListBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder getAppendListBuilder() { + return getAppendListFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder getAppendListOrBuilder() { + if ((methodCase_ == 6) && (appendListBuilder_ != null)) { + return appendListBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 6) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.AppendList appendList = 6; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder> + getAppendListFieldBuilder() { + if (appendListBuilder_ == null) { + if (!(methodCase_ == 6)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); + } + appendListBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 6; + onChanged();; + return appendListBuilder_; + } + + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder> clearBuilder_; + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return Whether the clear field is set. + */ + @java.lang.Override + public boolean hasClear() { + return methodCase_ == 7; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + * @return The clear. + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getClear() { + if (clearBuilder_ == null) { + if (methodCase_ == 7) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } else { + if (methodCase_ == 7) { + return clearBuilder_.getMessage(); + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + public Builder setClear(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear value) { + if (clearBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + method_ = value; + onChanged(); + } else { + clearBuilder_.setMessage(value); + } + methodCase_ = 7; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + public Builder setClear( + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder builderForValue) { + if (clearBuilder_ == null) { + method_ = builderForValue.build(); + onChanged(); + } else { + clearBuilder_.setMessage(builderForValue.build()); + } + methodCase_ = 7; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + public Builder mergeClear(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear value) { + if (clearBuilder_ == null) { + if (methodCase_ == 7 && + method_ != org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance()) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.newBuilder((org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_) + .mergeFrom(value).buildPartial(); + } else { + method_ = value; + } + onChanged(); + } else { + if (methodCase_ == 7) { + clearBuilder_.mergeFrom(value); + } else { + clearBuilder_.setMessage(value); + } + } + methodCase_ = 7; + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + public Builder clearClear() { + if (clearBuilder_ == null) { + if (methodCase_ == 7) { + methodCase_ = 0; + method_ = null; + onChanged(); + } + } else { + if (methodCase_ == 7) { + methodCase_ = 0; + method_ = null; + } + clearBuilder_.clear(); + } + return this; + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder getClearBuilder() { + return getClearFieldBuilder().getBuilder(); + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder getClearOrBuilder() { + if ((methodCase_ == 7) && (clearBuilder_ != null)) { + return clearBuilder_.getMessageOrBuilder(); + } else { + if (methodCase_ == 7) { + return (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_; + } + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } + } + /** + * .org.apache.spark.sql.execution.streaming.state.Clear clear = 7; + */ + private com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder> + getClearFieldBuilder() { + if (clearBuilder_ == null) { + if (!(methodCase_ == 7)) { + method_ = org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + } + clearBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder, org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder>( + (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) method_, + getParentForChildren(), + isClean()); + method_ = null; + } + methodCase_ = 7; + onChanged();; + return clearBuilder_; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.ListStateCall) + } + + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.ListStateCall) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall(); + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public ListStateCall parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateCall getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface SetImplicitKeyOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + com.google.protobuf.MessageOrBuilder { + + /** + * bytes key = 1; + * @return The key. + */ + com.google.protobuf.ByteString getKey(); + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.SetImplicitKey} + */ + public static final class SetImplicitKey extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + SetImplicitKeyOrBuilder { + private static final long serialVersionUID = 0L; + // Use SetImplicitKey.newBuilder() to construct. + private SetImplicitKey(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private SetImplicitKey() { + key_ = com.google.protobuf.ByteString.EMPTY; + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new SetImplicitKey(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.Builder.class); + } + + public static final int KEY_FIELD_NUMBER = 1; + private com.google.protobuf.ByteString key_; + /** + * bytes key = 1; + * @return The key. + */ + @java.lang.Override + public com.google.protobuf.ByteString getKey() { + return key_; + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + if (!key_.isEmpty()) { + output.writeBytes(1, key_); + } + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + if (!key_.isEmpty()) { + size += com.google.protobuf.CodedOutputStream + .computeBytesSize(1, key_); + } + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey)) { + return super.equals(obj); + } + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey other = (org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey) obj; + + if (!getKey() + .equals(other.getKey())) return false; + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (37 * hash) + KEY_FIELD_NUMBER; + hash = (53 * hash) + getKey().hashCode(); + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.SetImplicitKey} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKeyOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.Builder.class); + } + + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + key_ = com.google.protobuf.ByteString.EMPTY; + + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.getDefaultInstance(); + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey result = new org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey(this); + result.key_ = key_; + onBuilt(); + return result; + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.getDefaultInstance()) return this; + if (other.getKey() != com.google.protobuf.ByteString.EMPTY) { + setKey(other.getKey()); + } + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + case 10: { + key_ = input.readBytes(); + + break; + } // case 10 + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + + private com.google.protobuf.ByteString key_ = com.google.protobuf.ByteString.EMPTY; + /** + * bytes key = 1; + * @return The key. + */ + @java.lang.Override + public com.google.protobuf.ByteString getKey() { + return key_; + } + /** + * bytes key = 1; + * @param value The key to set. + * @return This builder for chaining. + */ + public Builder setKey(com.google.protobuf.ByteString value) { + if (value == null) { + throw new NullPointerException(); + } + + key_ = value; + onChanged(); + return this; + } + /** + * bytes key = 1; + * @return This builder for chaining. + */ + public Builder clearKey() { + + key_ = getDefaultInstance().getKey(); + onChanged(); + return this; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + } + + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey(); + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public SetImplicitKey parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface RemoveImplicitKeyOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + com.google.protobuf.MessageOrBuilder { + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey} + */ + public static final class RemoveImplicitKey extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + RemoveImplicitKeyOrBuilder { + private static final long serialVersionUID = 0L; + // Use RemoveImplicitKey.newBuilder() to construct. + private RemoveImplicitKey(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private RemoveImplicitKey() { + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new RemoveImplicitKey(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.Builder.class); + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey)) { + return super.equals(obj); + } + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey other = (org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey) obj; + + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKeyOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.Builder.class); + } + + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.getDefaultInstance(); + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey result = new org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey(this); + onBuilt(); + return result; + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.getDefaultInstance()) return this; + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + } + + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey(); + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public RemoveImplicitKey parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface ExistsOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Exists) + com.google.protobuf.MessageOrBuilder { + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Exists} + */ + public static final class Exists extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Exists) + ExistsOrBuilder { + private static final long serialVersionUID = 0L; + // Use Exists.newBuilder() to construct. + private Exists(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private Exists() { + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new Exists(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder.class); + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Exists)) { + return super.equals(obj); + } + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) obj; + + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Exists} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Exists) + org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder.class); + } + + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Exists result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Exists(this); + onBuilt(); + return result; + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Exists)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance()) return this; + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Exists) + } + + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Exists) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Exists DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Exists(); + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public Exists parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface GetOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Get) + com.google.protobuf.MessageOrBuilder { + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Get} + */ + public static final class Get extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Get) + GetOrBuilder { + private static final long serialVersionUID = 0L; + // Use Get.newBuilder() to construct. + private Get(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private Get() { + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new Get(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.Get.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Get.Builder.class); + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Get)) { + return super.equals(obj); + } + org.apache.spark.sql.execution.streaming.state.StateMessage.Get other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Get) obj; + + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Get prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Get} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Get) + org.apache.spark.sql.execution.streaming.state.StateMessage.GetOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.Get.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Get.Builder.class); + } + + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Get.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.Get.getDefaultInstance(); + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Get build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Get result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Get buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Get result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Get(this); + onBuilt(); + return result; + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Get) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Get)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Get other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Get.getDefaultInstance()) return this; + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Get) + } + + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Get) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Get DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Get(); + } + + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public Get parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface ValueStateUpdateOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + com.google.protobuf.MessageOrBuilder { + + /** + * bytes value = 1; + * @return The value. + */ + com.google.protobuf.ByteString getValue(); + } + /** + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ValueStateUpdate} + */ + public static final class ValueStateUpdate extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + ValueStateUpdateOrBuilder { + private static final long serialVersionUID = 0L; + // Use ValueStateUpdate.newBuilder() to construct. + private ValueStateUpdate(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private ValueStateUpdate() { + value_ = com.google.protobuf.ByteString.EMPTY; + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new ValueStateUpdate(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_fieldAccessorTable + .ensureFieldAccessorsInitialized( + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.Builder.class); + } + + public static final int VALUE_FIELD_NUMBER = 1; + private com.google.protobuf.ByteString value_; + /** + * bytes value = 1; + * @return The value. + */ + @java.lang.Override + public com.google.protobuf.ByteString getValue() { + return value_; + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { byte isInitialized = memoizedIsInitialized; if (isInitialized == 1) return true; if (isInitialized == 0) return false; @@ -7558,8 +11434,8 @@ public final boolean isInitialized() { @java.lang.Override public void writeTo(com.google.protobuf.CodedOutputStream output) throws java.io.IOException { - if (!key_.isEmpty()) { - output.writeBytes(1, key_); + if (!value_.isEmpty()) { + output.writeBytes(1, value_); } getUnknownFields().writeTo(output); } @@ -7570,9 +11446,9 @@ public int getSerializedSize() { if (size != -1) return size; size = 0; - if (!key_.isEmpty()) { + if (!value_.isEmpty()) { size += com.google.protobuf.CodedOutputStream - .computeBytesSize(1, key_); + .computeBytesSize(1, value_); } size += getUnknownFields().getSerializedSize(); memoizedSize = size; @@ -7584,13 +11460,13 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey other = (org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate other = (org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate) obj; - if (!getKey() - .equals(other.getKey())) return false; + if (!getValue() + .equals(other.getValue())) return false; if (!getUnknownFields().equals(other.getUnknownFields())) return false; return true; } @@ -7602,76 +11478,76 @@ public int hashCode() { } int hash = 41; hash = (19 * hash) + getDescriptor().hashCode(); - hash = (37 * hash) + KEY_FIELD_NUMBER; - hash = (53 * hash) + getKey().hashCode(); + hash = (37 * hash) + VALUE_FIELD_NUMBER; + hash = (53 * hash) + getValue().hashCode(); hash = (29 * hash) + getUnknownFields().hashCode(); memoizedHashCode = hash; return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -7684,7 +11560,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImp public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -7700,26 +11576,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.SetImplicitKey} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ValueStateUpdate} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKeyOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdateOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.newBuilder() private Builder() { } @@ -7732,7 +11608,7 @@ private Builder( @java.lang.Override public Builder clear() { super.clear(); - key_ = com.google.protobuf.ByteString.EMPTY; + value_ = com.google.protobuf.ByteString.EMPTY; return this; } @@ -7740,17 +11616,17 @@ public Builder clear() { @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -7758,9 +11634,9 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKe } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey result = new org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey(this); - result.key_ = key_; + public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate result = new org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate(this); + result.value_ = value_; onBuilt(); return result; } @@ -7799,18 +11675,18 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey.getDefaultInstance()) return this; - if (other.getKey() != com.google.protobuf.ByteString.EMPTY) { - setKey(other.getKey()); + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.getDefaultInstance()) return this; + if (other.getValue() != com.google.protobuf.ByteString.EMPTY) { + setValue(other.getValue()); } this.mergeUnknownFields(other.getUnknownFields()); onChanged(); @@ -7839,7 +11715,7 @@ public Builder mergeFrom( done = true; break; case 10: { - key_ = input.readBytes(); + value_ = input.readBytes(); break; } // case 10 @@ -7859,36 +11735,36 @@ public Builder mergeFrom( return this; } - private com.google.protobuf.ByteString key_ = com.google.protobuf.ByteString.EMPTY; + private com.google.protobuf.ByteString value_ = com.google.protobuf.ByteString.EMPTY; /** - * bytes key = 1; - * @return The key. + * bytes value = 1; + * @return The value. */ @java.lang.Override - public com.google.protobuf.ByteString getKey() { - return key_; + public com.google.protobuf.ByteString getValue() { + return value_; } /** - * bytes key = 1; - * @param value The key to set. + * bytes value = 1; + * @param value The value to set. * @return This builder for chaining. */ - public Builder setKey(com.google.protobuf.ByteString value) { + public Builder setValue(com.google.protobuf.ByteString value) { if (value == null) { throw new NullPointerException(); } - key_ = value; + value_ = value; onChanged(); return this; } /** - * bytes key = 1; + * bytes value = 1; * @return This builder for chaining. */ - public Builder clearKey() { + public Builder clearValue() { - key_ = getDefaultInstance().getKey(); + value_ = getDefaultInstance().getValue(); onChanged(); return this; } @@ -7905,23 +11781,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.SetImplicitKey) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public SetImplicitKey parsePartialFrom( + public ValueStateUpdate parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -7940,46 +11816,46 @@ public SetImplicitKey parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.SetImplicitKey getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstanceForType() { return DEFAULT_INSTANCE; } } - public interface RemoveImplicitKeyOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + public interface ClearOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Clear) com.google.protobuf.MessageOrBuilder { } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Clear} */ - public static final class RemoveImplicitKey extends + public static final class Clear extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) - RemoveImplicitKeyOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Clear) + ClearOrBuilder { private static final long serialVersionUID = 0L; - // Use RemoveImplicitKey.newBuilder() to construct. - private RemoveImplicitKey(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use Clear.newBuilder() to construct. + private Clear(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private RemoveImplicitKey() { + private Clear() { } @java.lang.Override @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new RemoveImplicitKey(); + return new Clear(); } @java.lang.Override @@ -7989,15 +11865,15 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder.class); } private byte memoizedIsInitialized = -1; @@ -8033,10 +11909,10 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Clear)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey other = (org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) obj; if (!getUnknownFields().equals(other.getUnknownFields())) return false; return true; @@ -8054,69 +11930,69 @@ public int hashCode() { return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -8129,7 +12005,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.Remove public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -8145,26 +12021,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Clear} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKeyOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Clear) + org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.class, org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.newBuilder() private Builder() { } @@ -8183,17 +12059,17 @@ public Builder clear() { @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -8201,8 +12077,8 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplici } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey result = new org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey(this); + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.Clear result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Clear(this); onBuilt(); return result; } @@ -8241,16 +12117,16 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Clear)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey.getDefaultInstance()) return this; + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance()) return this; this.mergeUnknownFields(other.getUnknownFields()); onChanged(); return this; @@ -8305,23 +12181,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Clear) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.RemoveImplicitKey) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Clear) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Clear DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Clear(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public RemoveImplicitKey parsePartialFrom( + public Clear parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -8340,46 +12216,59 @@ public RemoveImplicitKey parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.RemoveImplicitKey getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstanceForType() { return DEFAULT_INSTANCE; } } - public interface ExistsOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Exists) + public interface ListStateGetOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.ListStateGet) com.google.protobuf.MessageOrBuilder { + + /** + * string iteratorId = 1; + * @return The iteratorId. + */ + java.lang.String getIteratorId(); + /** + * string iteratorId = 1; + * @return The bytes for iteratorId. + */ + com.google.protobuf.ByteString + getIteratorIdBytes(); } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Exists} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStateGet} */ - public static final class Exists extends + public static final class ListStateGet extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Exists) - ExistsOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.ListStateGet) + ListStateGetOrBuilder { private static final long serialVersionUID = 0L; - // Use Exists.newBuilder() to construct. - private Exists(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use ListStateGet.newBuilder() to construct. + private ListStateGet(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private Exists() { + private ListStateGet() { + iteratorId_ = ""; } @java.lang.Override @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new Exists(); + return new ListStateGet(); } @java.lang.Override @@ -8389,15 +12278,53 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder.class); + } + + public static final int ITERATORID_FIELD_NUMBER = 1; + private volatile java.lang.Object iteratorId_; + /** + * string iteratorId = 1; + * @return The iteratorId. + */ + @java.lang.Override + public java.lang.String getIteratorId() { + java.lang.Object ref = iteratorId_; + if (ref instanceof java.lang.String) { + return (java.lang.String) ref; + } else { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + iteratorId_ = s; + return s; + } + } + /** + * string iteratorId = 1; + * @return The bytes for iteratorId. + */ + @java.lang.Override + public com.google.protobuf.ByteString + getIteratorIdBytes() { + java.lang.Object ref = iteratorId_; + if (ref instanceof java.lang.String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + iteratorId_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } } private byte memoizedIsInitialized = -1; @@ -8414,6 +12341,9 @@ public final boolean isInitialized() { @java.lang.Override public void writeTo(com.google.protobuf.CodedOutputStream output) throws java.io.IOException { + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(iteratorId_)) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 1, iteratorId_); + } getUnknownFields().writeTo(output); } @@ -8423,6 +12353,9 @@ public int getSerializedSize() { if (size != -1) return size; size = 0; + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(iteratorId_)) { + size += com.google.protobuf.GeneratedMessageV3.computeStringSize(1, iteratorId_); + } size += getUnknownFields().getSerializedSize(); memoizedSize = size; return size; @@ -8433,11 +12366,13 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Exists)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.Exists other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet other = (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) obj; + if (!getIteratorId() + .equals(other.getIteratorId())) return false; if (!getUnknownFields().equals(other.getUnknownFields())) return false; return true; } @@ -8449,74 +12384,76 @@ public int hashCode() { } int hash = 41; hash = (19 * hash) + getDescriptor().hashCode(); + hash = (37 * hash) + ITERATORID_FIELD_NUMBER; + hash = (53 * hash) + getIteratorId().hashCode(); hash = (29 * hash) + getUnknownFields().hashCode(); memoizedHashCode = hash; return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -8529,7 +12466,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -8545,26 +12482,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Exists} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStateGet} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Exists) - org.apache.spark.sql.execution.streaming.state.StateMessage.ExistsOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.ListStateGet) + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGetOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.newBuilder() private Builder() { } @@ -8577,23 +12514,25 @@ private Builder( @java.lang.Override public Builder clear() { super.clear(); + iteratorId_ = ""; + return this; } @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Exists result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -8601,8 +12540,9 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists build( } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Exists result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Exists(this); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet result = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet(this); + result.iteratorId_ = iteratorId_; onBuilt(); return result; } @@ -8641,16 +12581,20 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Exists) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Exists)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Exists other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Exists.getDefaultInstance()) return this; + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet.getDefaultInstance()) return this; + if (!other.getIteratorId().isEmpty()) { + iteratorId_ = other.iteratorId_; + onChanged(); + } this.mergeUnknownFields(other.getUnknownFields()); onChanged(); return this; @@ -8677,6 +12621,11 @@ public Builder mergeFrom( case 0: done = true; break; + case 10: { + iteratorId_ = input.readStringRequireUtf8(); + + break; + } // case 10 default: { if (!super.parseUnknownField(input, extensionRegistry, tag)) { done = true; // was an endgroup tag @@ -8692,6 +12641,82 @@ public Builder mergeFrom( } // finally return this; } + + private java.lang.Object iteratorId_ = ""; + /** + * string iteratorId = 1; + * @return The iteratorId. + */ + public java.lang.String getIteratorId() { + java.lang.Object ref = iteratorId_; + if (!(ref instanceof java.lang.String)) { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + iteratorId_ = s; + return s; + } else { + return (java.lang.String) ref; + } + } + /** + * string iteratorId = 1; + * @return The bytes for iteratorId. + */ + public com.google.protobuf.ByteString + getIteratorIdBytes() { + java.lang.Object ref = iteratorId_; + if (ref instanceof String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + iteratorId_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + /** + * string iteratorId = 1; + * @param value The iteratorId to set. + * @return This builder for chaining. + */ + public Builder setIteratorId( + java.lang.String value) { + if (value == null) { + throw new NullPointerException(); + } + + iteratorId_ = value; + onChanged(); + return this; + } + /** + * string iteratorId = 1; + * @return This builder for chaining. + */ + public Builder clearIteratorId() { + + iteratorId_ = getDefaultInstance().getIteratorId(); + onChanged(); + return this; + } + /** + * string iteratorId = 1; + * @param value The bytes for iteratorId to set. + * @return This builder for chaining. + */ + public Builder setIteratorIdBytes( + com.google.protobuf.ByteString value) { + if (value == null) { + throw new NullPointerException(); + } + checkByteStringIsUtf8(value); + + iteratorId_ = value; + onChanged(); + return this; + } @java.lang.Override public final Builder setUnknownFields( final com.google.protobuf.UnknownFieldSet unknownFields) { @@ -8705,23 +12730,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Exists) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.ListStateGet) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Exists) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Exists DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.ListStateGet) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Exists(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public Exists parsePartialFrom( + public ListStateGet parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -8740,46 +12765,46 @@ public Exists parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Exists getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStateGet getDefaultInstanceForType() { return DEFAULT_INSTANCE; } } - public interface GetOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Get) + public interface ListStatePutOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.ListStatePut) com.google.protobuf.MessageOrBuilder { } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Get} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStatePut} */ - public static final class Get extends + public static final class ListStatePut extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Get) - GetOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.ListStatePut) + ListStatePutOrBuilder { private static final long serialVersionUID = 0L; - // Use Get.newBuilder() to construct. - private Get(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use ListStatePut.newBuilder() to construct. + private ListStatePut(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private Get() { + private ListStatePut() { } @java.lang.Override @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new Get(); + return new ListStatePut(); } @java.lang.Override @@ -8789,15 +12814,15 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Get.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Get.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder.class); } private byte memoizedIsInitialized = -1; @@ -8833,10 +12858,10 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Get)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.Get other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Get) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut other = (org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) obj; if (!getUnknownFields().equals(other.getUnknownFields())) return false; return true; @@ -8854,69 +12879,69 @@ public int hashCode() { return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -8929,7 +12954,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get pa public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Get prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -8945,26 +12970,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Get} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ListStatePut} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Get) - org.apache.spark.sql.execution.streaming.state.StateMessage.GetOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.ListStatePut) + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePutOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Get.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Get.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Get.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.newBuilder() private Builder() { } @@ -8983,17 +13008,17 @@ public Builder clear() { @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.Get.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Get build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Get result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -9001,8 +13026,8 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.Get build() { } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Get buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Get result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Get(this); + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut result = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut(this); onBuilt(); return result; } @@ -9041,16 +13066,16 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Get) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Get)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Get other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Get.getDefaultInstance()) return this; + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut.getDefaultInstance()) return this; this.mergeUnknownFields(other.getUnknownFields()); onChanged(); return this; @@ -9105,23 +13130,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Get) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.ListStatePut) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Get) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Get DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.ListStatePut) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Get(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public Get parsePartialFrom( + public ListStatePut parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -9140,24 +13165,24 @@ public Get parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Get getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.ListStatePut getDefaultInstanceForType() { return DEFAULT_INSTANCE; } } - public interface ValueStateUpdateOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + public interface AppendValueOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.AppendValue) com.google.protobuf.MessageOrBuilder { /** @@ -9167,18 +13192,18 @@ public interface ValueStateUpdateOrBuilder extends com.google.protobuf.ByteString getValue(); } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ValueStateUpdate} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.AppendValue} */ - public static final class ValueStateUpdate extends + public static final class AppendValue extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) - ValueStateUpdateOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.AppendValue) + AppendValueOrBuilder { private static final long serialVersionUID = 0L; - // Use ValueStateUpdate.newBuilder() to construct. - private ValueStateUpdate(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use AppendValue.newBuilder() to construct. + private AppendValue(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private ValueStateUpdate() { + private AppendValue() { value_ = com.google.protobuf.ByteString.EMPTY; } @@ -9186,7 +13211,7 @@ private ValueStateUpdate() { @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new ValueStateUpdate(); + return new AppendValue(); } @java.lang.Override @@ -9196,15 +13221,15 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.class, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder.class); } public static final int VALUE_FIELD_NUMBER = 1; @@ -9258,10 +13283,10 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate other = (org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue other = (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) obj; if (!getValue() .equals(other.getValue())) return false; @@ -9283,69 +13308,69 @@ public int hashCode() { return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -9358,7 +13383,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueS public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -9374,26 +13399,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.ValueStateUpdate} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.AppendValue} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdateOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.AppendValue) + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValueOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.class, org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.class, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.newBuilder() private Builder() { } @@ -9414,17 +13439,17 @@ public Builder clear() { @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -9432,8 +13457,8 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpd } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate result = new org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate(this); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue result = new org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue(this); result.value_ = value_; onBuilt(); return result; @@ -9473,16 +13498,16 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate.getDefaultInstance()) return this; + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue.getDefaultInstance()) return this; if (other.getValue() != com.google.protobuf.ByteString.EMPTY) { setValue(other.getValue()); } @@ -9579,23 +13604,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.AppendValue) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.ValueStateUpdate) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.AppendValue) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public ValueStateUpdate parsePartialFrom( + public AppendValue parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -9614,46 +13639,46 @@ public ValueStateUpdate parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.ValueStateUpdate getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendValue getDefaultInstanceForType() { return DEFAULT_INSTANCE; } } - public interface ClearOrBuilder extends - // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.Clear) + public interface AppendListOrBuilder extends + // @@protoc_insertion_point(interface_extends:org.apache.spark.sql.execution.streaming.state.AppendList) com.google.protobuf.MessageOrBuilder { } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Clear} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.AppendList} */ - public static final class Clear extends + public static final class AppendList extends com.google.protobuf.GeneratedMessageV3 implements - // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.Clear) - ClearOrBuilder { + // @@protoc_insertion_point(message_implements:org.apache.spark.sql.execution.streaming.state.AppendList) + AppendListOrBuilder { private static final long serialVersionUID = 0L; - // Use Clear.newBuilder() to construct. - private Clear(com.google.protobuf.GeneratedMessageV3.Builder builder) { + // Use AppendList.newBuilder() to construct. + private AppendList(com.google.protobuf.GeneratedMessageV3.Builder builder) { super(builder); } - private Clear() { + private AppendList() { } @java.lang.Override @SuppressWarnings({"unused"}) protected java.lang.Object newInstance( UnusedPrivateParameter unused) { - return new Clear(); + return new AppendList(); } @java.lang.Override @@ -9663,15 +13688,15 @@ protected java.lang.Object newInstance( } public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.class, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder.class); } private byte memoizedIsInitialized = -1; @@ -9707,10 +13732,10 @@ public boolean equals(final java.lang.Object obj) { if (obj == this) { return true; } - if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Clear)) { + if (!(obj instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList)) { return super.equals(obj); } - org.apache.spark.sql.execution.streaming.state.StateMessage.Clear other = (org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) obj; + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList other = (org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) obj; if (!getUnknownFields().equals(other.getUnknownFields())) return false; return true; @@ -9728,69 +13753,69 @@ public int hashCode() { return hash; } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( java.nio.ByteBuffer data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( java.nio.ByteBuffer data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( com.google.protobuf.ByteString data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( com.google.protobuf.ByteString data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom(byte[] data) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom(byte[] data) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( byte[] data, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { return PARSER.parseFrom(data, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseDelimitedFrom(java.io.InputStream input) + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseDelimitedFrom(java.io.InputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseDelimitedFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseDelimitedFrom( java.io.InputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseDelimitedWithIOException(PARSER, input, extensionRegistry); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( com.google.protobuf.CodedInputStream input) throws java.io.IOException { return com.google.protobuf.GeneratedMessageV3 .parseWithIOException(PARSER, input); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear parseFrom( + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList parseFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws java.io.IOException { @@ -9803,7 +13828,7 @@ public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear public static Builder newBuilder() { return DEFAULT_INSTANCE.toBuilder(); } - public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear prototype) { + public static Builder newBuilder(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList prototype) { return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); } @java.lang.Override @@ -9819,26 +13844,26 @@ protected Builder newBuilderForType( return builder; } /** - * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.Clear} + * Protobuf type {@code org.apache.spark.sql.execution.streaming.state.AppendList} */ public static final class Builder extends com.google.protobuf.GeneratedMessageV3.Builder implements - // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.Clear) - org.apache.spark.sql.execution.streaming.state.StateMessage.ClearOrBuilder { + // @@protoc_insertion_point(builder_implements:org.apache.spark.sql.execution.streaming.state.AppendList) + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendListOrBuilder { public static final com.google.protobuf.Descriptors.Descriptor getDescriptor() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor; } @java.lang.Override protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internalGetFieldAccessorTable() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_fieldAccessorTable .ensureFieldAccessorsInitialized( - org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.class, org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.Builder.class); + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.class, org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.Builder.class); } - // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.newBuilder() + // Construct using org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.newBuilder() private Builder() { } @@ -9857,17 +13882,17 @@ public Builder clear() { @java.lang.Override public com.google.protobuf.Descriptors.Descriptor getDescriptorForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor; + return org.apache.spark.sql.execution.streaming.state.StateMessage.internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstanceForType() { - return org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getDefaultInstanceForType() { + return org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance(); } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear build() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Clear result = buildPartial(); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList build() { + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList result = buildPartial(); if (!result.isInitialized()) { throw newUninitializedMessageException(result); } @@ -9875,8 +13900,8 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear build() } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear buildPartial() { - org.apache.spark.sql.execution.streaming.state.StateMessage.Clear result = new org.apache.spark.sql.execution.streaming.state.StateMessage.Clear(this); + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList buildPartial() { + org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList result = new org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList(this); onBuilt(); return result; } @@ -9915,16 +13940,16 @@ public Builder addRepeatedField( } @java.lang.Override public Builder mergeFrom(com.google.protobuf.Message other) { - if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.Clear) { - return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.Clear)other); + if (other instanceof org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList) { + return mergeFrom((org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList)other); } else { super.mergeFrom(other); return this; } } - public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.Clear other) { - if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.Clear.getDefaultInstance()) return this; + public Builder mergeFrom(org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList other) { + if (other == org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList.getDefaultInstance()) return this; this.mergeUnknownFields(other.getUnknownFields()); onChanged(); return this; @@ -9979,23 +14004,23 @@ public final Builder mergeUnknownFields( } - // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.Clear) + // @@protoc_insertion_point(builder_scope:org.apache.spark.sql.execution.streaming.state.AppendList) } - // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.Clear) - private static final org.apache.spark.sql.execution.streaming.state.StateMessage.Clear DEFAULT_INSTANCE; + // @@protoc_insertion_point(class_scope:org.apache.spark.sql.execution.streaming.state.AppendList) + private static final org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList DEFAULT_INSTANCE; static { - DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.Clear(); + DEFAULT_INSTANCE = new org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList(); } - public static org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstance() { + public static org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getDefaultInstance() { return DEFAULT_INSTANCE; } - private static final com.google.protobuf.Parser - PARSER = new com.google.protobuf.AbstractParser() { + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { @java.lang.Override - public Clear parsePartialFrom( + public AppendList parsePartialFrom( com.google.protobuf.CodedInputStream input, com.google.protobuf.ExtensionRegistryLite extensionRegistry) throws com.google.protobuf.InvalidProtocolBufferException { @@ -10014,17 +14039,17 @@ public Clear parsePartialFrom( } }; - public static com.google.protobuf.Parser parser() { + public static com.google.protobuf.Parser parser() { return PARSER; } @java.lang.Override - public com.google.protobuf.Parser getParserForType() { + public com.google.protobuf.Parser getParserForType() { return PARSER; } @java.lang.Override - public org.apache.spark.sql.execution.streaming.state.StateMessage.Clear getDefaultInstanceForType() { + public org.apache.spark.sql.execution.streaming.state.StateMessage.AppendList getDefaultInstanceForType() { return DEFAULT_INSTANCE; } @@ -11041,6 +15066,11 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.TTLConfig get private static final com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateCall_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_fieldAccessorTable; private static final com.google.protobuf.Descriptors.Descriptor internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor; private static final @@ -11071,6 +15101,26 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.TTLConfig get private static final com.google.protobuf.GeneratedMessageV3.FieldAccessorTable internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_fieldAccessorTable; private static final com.google.protobuf.Descriptors.Descriptor internal_static_org_apache_spark_sql_execution_streaming_state_SetHandleState_descriptor; private static final @@ -11112,36 +15162,54 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.TTLConfig get "xecution.streaming.state.StateCallComman" + "dH\000\022W\n\013getMapState\030\004 \001(\0132@.org.apache.sp" + "ark.sql.execution.streaming.state.StateC" + - "allCommandH\000B\010\n\006method\"z\n\024StateVariableR" + - "equest\022X\n\016valueStateCall\030\001 \001(\0132>.org.apa" + - "che.spark.sql.execution.streaming.state." + - "ValueStateCallH\000B\010\n\006method\"\340\001\n\032ImplicitG" + - "roupingKeyRequest\022X\n\016setImplicitKey\030\001 \001(" + - "\0132>.org.apache.spark.sql.execution.strea" + - "ming.state.SetImplicitKeyH\000\022^\n\021removeImp" + - "licitKey\030\002 \001(\0132A.org.apache.spark.sql.ex" + - "ecution.streaming.state.RemoveImplicitKe" + - "yH\000B\010\n\006method\"}\n\020StateCallCommand\022\021\n\tsta" + - "teName\030\001 \001(\t\022\016\n\006schema\030\002 \001(\t\022F\n\003ttl\030\003 \001(" + - "\01329.org.apache.spark.sql.execution.strea" + - "ming.state.TTLConfig\"\341\002\n\016ValueStateCall\022" + - "\021\n\tstateName\030\001 \001(\t\022H\n\006exists\030\002 \001(\01326.org" + + "allCommandH\000B\010\n\006method\"\322\001\n\024StateVariable" + + "Request\022X\n\016valueStateCall\030\001 \001(\0132>.org.ap" + + "ache.spark.sql.execution.streaming.state" + + ".ValueStateCallH\000\022V\n\rlistStateCall\030\002 \001(\013" + + "2=.org.apache.spark.sql.execution.stream" + + "ing.state.ListStateCallH\000B\010\n\006method\"\340\001\n\032" + + "ImplicitGroupingKeyRequest\022X\n\016setImplici" + + "tKey\030\001 \001(\0132>.org.apache.spark.sql.execut" + + "ion.streaming.state.SetImplicitKeyH\000\022^\n\021" + + "removeImplicitKey\030\002 \001(\0132A.org.apache.spa" + + "rk.sql.execution.streaming.state.RemoveI" + + "mplicitKeyH\000B\010\n\006method\"}\n\020StateCallComma" + + "nd\022\021\n\tstateName\030\001 \001(\t\022\016\n\006schema\030\002 \001(\t\022F\n" + + "\003ttl\030\003 \001(\01329.org.apache.spark.sql.execut" + + "ion.streaming.state.TTLConfig\"\341\002\n\016ValueS" + + "tateCall\022\021\n\tstateName\030\001 \001(\t\022H\n\006exists\030\002 " + + "\001(\01326.org.apache.spark.sql.execution.str" + + "eaming.state.ExistsH\000\022B\n\003get\030\003 \001(\01323.org" + ".apache.spark.sql.execution.streaming.st" + - "ate.ExistsH\000\022B\n\003get\030\003 \001(\01323.org.apache.s" + - "park.sql.execution.streaming.state.GetH\000" + - "\022\\\n\020valueStateUpdate\030\004 \001(\0132@.org.apache." + - "spark.sql.execution.streaming.state.Valu" + - "eStateUpdateH\000\022F\n\005clear\030\005 \001(\01325.org.apac" + - "he.spark.sql.execution.streaming.state.C" + - "learH\000B\010\n\006method\"\035\n\016SetImplicitKey\022\013\n\003ke" + - "y\030\001 \001(\014\"\023\n\021RemoveImplicitKey\"\010\n\006Exists\"\005" + - "\n\003Get\"!\n\020ValueStateUpdate\022\r\n\005value\030\001 \001(\014" + - "\"\007\n\005Clear\"\\\n\016SetHandleState\022J\n\005state\030\001 \001" + - "(\0162;.org.apache.spark.sql.execution.stre" + - "aming.state.HandleState\"\037\n\tTTLConfig\022\022\n\n" + - "durationMs\030\001 \001(\005*K\n\013HandleState\022\013\n\007CREAT" + - "ED\020\000\022\017\n\013INITIALIZED\020\001\022\022\n\016DATA_PROCESSED\020" + - "\002\022\n\n\006CLOSED\020\003b\006proto3" + "ate.GetH\000\022\\\n\020valueStateUpdate\030\004 \001(\0132@.or" + + "g.apache.spark.sql.execution.streaming.s" + + "tate.ValueStateUpdateH\000\022F\n\005clear\030\005 \001(\01325" + + ".org.apache.spark.sql.execution.streamin" + + "g.state.ClearH\000B\010\n\006method\"\220\004\n\rListStateC" + + "all\022\021\n\tstateName\030\001 \001(\t\022H\n\006exists\030\002 \001(\01326" + + ".org.apache.spark.sql.execution.streamin" + + "g.state.ExistsH\000\022T\n\014listStateGet\030\003 \001(\0132<" + + ".org.apache.spark.sql.execution.streamin" + + "g.state.ListStateGetH\000\022T\n\014listStatePut\030\004" + + " \001(\0132<.org.apache.spark.sql.execution.st" + + "reaming.state.ListStatePutH\000\022R\n\013appendVa" + + "lue\030\005 \001(\0132;.org.apache.spark.sql.executi" + + "on.streaming.state.AppendValueH\000\022P\n\nappe" + + "ndList\030\006 \001(\0132:.org.apache.spark.sql.exec" + + "ution.streaming.state.AppendListH\000\022F\n\005cl" + + "ear\030\007 \001(\01325.org.apache.spark.sql.executi" + + "on.streaming.state.ClearH\000B\010\n\006method\"\035\n\016" + + "SetImplicitKey\022\013\n\003key\030\001 \001(\014\"\023\n\021RemoveImp" + + "licitKey\"\010\n\006Exists\"\005\n\003Get\"!\n\020ValueStateU" + + "pdate\022\r\n\005value\030\001 \001(\014\"\007\n\005Clear\"\"\n\014ListSta" + + "teGet\022\022\n\niteratorId\030\001 \001(\t\"\016\n\014ListStatePu" + + "t\"\034\n\013AppendValue\022\r\n\005value\030\001 \001(\014\"\014\n\nAppen" + + "dList\"\\\n\016SetHandleState\022J\n\005state\030\001 \001(\0162;" + + ".org.apache.spark.sql.execution.streamin" + + "g.state.HandleState\"\037\n\tTTLConfig\022\022\n\ndura" + + "tionMs\030\001 \001(\005*K\n\013HandleState\022\013\n\007CREATED\020\000" + + "\022\017\n\013INITIALIZED\020\001\022\022\n\016DATA_PROCESSED\020\002\022\n\n" + + "\006CLOSED\020\003b\006proto3" }; descriptor = com.google.protobuf.Descriptors.FileDescriptor .internalBuildGeneratedFileFrom(descriptorData, @@ -11170,7 +15238,7 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.TTLConfig get internal_static_org_apache_spark_sql_execution_streaming_state_StateVariableRequest_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_StateVariableRequest_descriptor, - new java.lang.String[] { "ValueStateCall", "Method", }); + new java.lang.String[] { "ValueStateCall", "ListStateCall", "Method", }); internal_static_org_apache_spark_sql_execution_streaming_state_ImplicitGroupingKeyRequest_descriptor = getDescriptor().getMessageTypes().get(4); internal_static_org_apache_spark_sql_execution_streaming_state_ImplicitGroupingKeyRequest_fieldAccessorTable = new @@ -11189,50 +15257,80 @@ public org.apache.spark.sql.execution.streaming.state.StateMessage.TTLConfig get com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateCall_descriptor, new java.lang.String[] { "StateName", "Exists", "Get", "ValueStateUpdate", "Clear", "Method", }); - internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor = + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor = getDescriptor().getMessageTypes().get(7); + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateCall_descriptor, + new java.lang.String[] { "StateName", "Exists", "ListStateGet", "ListStatePut", "AppendValue", "AppendList", "Clear", "Method", }); + internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor = + getDescriptor().getMessageTypes().get(8); internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_SetImplicitKey_descriptor, new java.lang.String[] { "Key", }); internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor = - getDescriptor().getMessageTypes().get(8); + getDescriptor().getMessageTypes().get(9); internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_RemoveImplicitKey_descriptor, new java.lang.String[] { }); internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor = - getDescriptor().getMessageTypes().get(9); + getDescriptor().getMessageTypes().get(10); internal_static_org_apache_spark_sql_execution_streaming_state_Exists_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_Exists_descriptor, new java.lang.String[] { }); internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor = - getDescriptor().getMessageTypes().get(10); + getDescriptor().getMessageTypes().get(11); internal_static_org_apache_spark_sql_execution_streaming_state_Get_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_Get_descriptor, new java.lang.String[] { }); internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor = - getDescriptor().getMessageTypes().get(11); + getDescriptor().getMessageTypes().get(12); internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_ValueStateUpdate_descriptor, new java.lang.String[] { "Value", }); internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor = - getDescriptor().getMessageTypes().get(12); + getDescriptor().getMessageTypes().get(13); internal_static_org_apache_spark_sql_execution_streaming_state_Clear_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_Clear_descriptor, new java.lang.String[] { }); + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor = + getDescriptor().getMessageTypes().get(14); + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_org_apache_spark_sql_execution_streaming_state_ListStateGet_descriptor, + new java.lang.String[] { "IteratorId", }); + internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor = + getDescriptor().getMessageTypes().get(15); + internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_org_apache_spark_sql_execution_streaming_state_ListStatePut_descriptor, + new java.lang.String[] { }); + internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor = + getDescriptor().getMessageTypes().get(16); + internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_org_apache_spark_sql_execution_streaming_state_AppendValue_descriptor, + new java.lang.String[] { "Value", }); + internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor = + getDescriptor().getMessageTypes().get(17); + internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_org_apache_spark_sql_execution_streaming_state_AppendList_descriptor, + new java.lang.String[] { }); internal_static_org_apache_spark_sql_execution_streaming_state_SetHandleState_descriptor = - getDescriptor().getMessageTypes().get(13); + getDescriptor().getMessageTypes().get(18); internal_static_org_apache_spark_sql_execution_streaming_state_SetHandleState_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_SetHandleState_descriptor, new java.lang.String[] { "State", }); internal_static_org_apache_spark_sql_execution_streaming_state_TTLConfig_descriptor = - getDescriptor().getMessageTypes().get(14); + getDescriptor().getMessageTypes().get(19); internal_static_org_apache_spark_sql_execution_streaming_state_TTLConfig_fieldAccessorTable = new com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( internal_static_org_apache_spark_sql_execution_streaming_state_TTLConfig_descriptor, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasDeserializer.scala new file mode 100644 index 0000000000000..82d4978853cb6 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasDeserializer.scala @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.io.DataInputStream + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.vector.ipc.ArrowStreamReader + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} + +/** + * A helper class to deserialize state Arrow batches from the state socket in + * TransformWithStateInPandas. + */ +class TransformWithStateInPandasDeserializer(deserializer: ExpressionEncoder.Deserializer[Row]) + extends Logging { + private val allocator = ArrowUtils.rootAllocator.newChildAllocator( + s"stdin reader for transformWithStateInPandas state socket", 0, Long.MaxValue) + + /** + * Read Arrow batches from the given stream and deserialize them into rows. + */ + def readArrowBatches(stream: DataInputStream): Seq[Row] = { + val reader = new ArrowStreamReader(stream, allocator) + val root = reader.getVectorSchemaRoot + val vectors = root.getFieldVectors.asScala.map { vector => + new ArrowColumnVector(vector) + }.toArray[ColumnVector] + val rows = ArrayBuffer[Row]() + while (reader.loadNextBatch()) { + val batch = new ColumnarBatch(vectors) + batch.setNumRows(root.getRowCount) + rows.appendAll(batch.rowIterator().asScala.map(r => deserializer(r.copy()))) + } + reader.close(false) + rows.toSeq + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala index 7d0c177d1df8f..b4b516ba9e5a6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala @@ -103,7 +103,8 @@ class TransformWithStateInPandasPythonRunner( executionContext.execute( new TransformWithStateInPandasStateServer(stateServerSocket, processorHandle, - groupingKeySchema)) + groupingKeySchema, timeZoneId, errorOnDuplicatedFieldNames, largeVarTypes, + sqlConf.arrowTransformWithStateInPandasMaxRecordsPerBatch)) context.addTaskCompletionListener[Unit] { _ => logInfo(log"completion listener called") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala index b5ec26b401d28..d293e7a4a5bb2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala @@ -24,15 +24,18 @@ import java.time.Duration import scala.collection.mutable import com.google.protobuf.ByteString +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.ArrowStreamWriter import org.apache.spark.internal.{Logging, LogKeys, MDC} import org.apache.spark.sql.{Encoders, Row} import org.apache.spark.sql.api.python.PythonSQLUtils import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder -import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, StatefulProcessorHandleImpl, StatefulProcessorHandleState} -import org.apache.spark.sql.execution.streaming.state.StateMessage.{HandleState, ImplicitGroupingKeyRequest, StatefulProcessorCall, StateRequest, StateResponse, StateVariableRequest, ValueStateCall} -import org.apache.spark.sql.streaming.{TTLConfig, ValueState} +import org.apache.spark.sql.execution.streaming.{ImplicitGroupingKeyTracker, StatefulProcessorHandleImpl, StatefulProcessorHandleState, StateVariableType} +import org.apache.spark.sql.execution.streaming.state.StateMessage.{HandleState, ImplicitGroupingKeyRequest, ListStateCall, StatefulProcessorCall, StateRequest, StateResponse, StateVariableRequest, ValueStateCall} +import org.apache.spark.sql.streaming.{ListState, TTLConfig, ValueState} import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.ArrowUtils /** * This class is used to handle the state requests from the Python side. It runs on a separate @@ -48,9 +51,16 @@ class TransformWithStateInPandasStateServer( stateServerSocket: ServerSocket, statefulProcessorHandle: StatefulProcessorHandleImpl, groupingKeySchema: StructType, + timeZoneId: String, + errorOnDuplicatedFieldNames: Boolean, + largeVarTypes: Boolean, + arrowTransformWithStateInPandasMaxRecordsPerBatch: Int, outputStreamForTest: DataOutputStream = null, - valueStateMapForTest: mutable.HashMap[String, - (ValueState[Row], StructType, ExpressionEncoder.Deserializer[Row])] = null) + valueStateMapForTest: mutable.HashMap[String, ValueStateInfo] = null, + deserializerForTest: TransformWithStateInPandasDeserializer = null, + arrowStreamWriterForTest: BaseStreamingArrowWriter = null, + listStatesMapForTest : mutable.HashMap[String, ListStateInfo] = null, + listStateIteratorMapForTest: mutable.HashMap[String, Iterator[Row]] = null) extends Runnable with Logging { private val keyRowDeserializer: ExpressionEncoder.Deserializer[Row] = ExpressionEncoder(groupingKeySchema).resolveAndBind().createDeserializer() @@ -60,8 +70,22 @@ class TransformWithStateInPandasStateServer( private val valueStates = if (valueStateMapForTest != null) { valueStateMapForTest } else { - new mutable.HashMap[String, (ValueState[Row], StructType, - ExpressionEncoder.Deserializer[Row])]() + new mutable.HashMap[String, ValueStateInfo]() + } + // A map to store the list state name -> (list state, schema, list state row deserializer, + // list state row serializer) mapping. + private val listStates = if (listStatesMapForTest != null) { + listStatesMapForTest + } else { + new mutable.HashMap[String, ListStateInfo]() + } + // A map to store the iterator id -> iterator mapping. This is to keep track of the + // current iterator position for each list state in a grouping key in case user tries to fetch + // another list state before the current iterator is exhausted. + private var listStateIterators = if (listStateIteratorMapForTest != null) { + listStateIteratorMapForTest + } else { + new mutable.HashMap[String, Iterator[Row]]() } def run(): Unit = { @@ -125,9 +149,13 @@ class TransformWithStateInPandasStateServer( // The key row is serialized as a byte array, we need to convert it back to a Row val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer) ImplicitGroupingKeyTracker.setImplicitKey(keyRow) + // Reset the list state iterators for a new grouping key. + listStateIterators = new mutable.HashMap[String, Iterator[Row]]() sendResponse(0) case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY => ImplicitGroupingKeyTracker.removeImplicitKey() + // Reset the list state iterators for a new grouping key. + listStateIterators = new mutable.HashMap[String, Iterator[Row]]() sendResponse(0) case _ => throw new IllegalArgumentException("Invalid method call") @@ -157,7 +185,12 @@ class TransformWithStateInPandasStateServer( val ttlDurationMs = if (message.getGetValueState.hasTtl) { Some(message.getGetValueState.getTtl.getDurationMs) } else None - initializeValueState(stateName, schema, ttlDurationMs) + initializeStateVariable(stateName, schema, StateVariableType.ValueState, ttlDurationMs) + case StatefulProcessorCall.MethodCase.GETLISTSTATE => + val stateName = message.getGetListState.getStateName + val schema = message.getGetListState.getSchema + // TODO(SPARK-49744): Add ttl support for list state. + initializeStateVariable(stateName, schema, StateVariableType.ListState, None) case _ => throw new IllegalArgumentException("Invalid method call") } @@ -167,6 +200,8 @@ class TransformWithStateInPandasStateServer( message.getMethodCase match { case StateVariableRequest.MethodCase.VALUESTATECALL => handleValueStateRequest(message.getValueStateCall) + case StateVariableRequest.MethodCase.LISTSTATECALL => + handleListStateRequest(message.getListStateCall) case _ => throw new IllegalArgumentException("Invalid method call") } @@ -179,16 +214,17 @@ class TransformWithStateInPandasStateServer( sendResponse(1, s"Value state $stateName is not initialized.") return } + val valueStateInfo = valueStates(stateName) message.getMethodCase match { case ValueStateCall.MethodCase.EXISTS => - if (valueStates(stateName)._1.exists()) { + if (valueStateInfo.valueState.exists()) { sendResponse(0) } else { // Send status code 2 to indicate that the value state doesn't have a value yet. sendResponse(2, s"state $stateName doesn't exist") } case ValueStateCall.MethodCase.GET => - val valueOption = valueStates(stateName)._1.getOption() + val valueOption = valueStateInfo.valueState.getOption() if (valueOption.isDefined) { // Serialize the value row as a byte array val valueBytes = PythonSQLUtils.toPyRow(valueOption.get) @@ -201,13 +237,95 @@ class TransformWithStateInPandasStateServer( } case ValueStateCall.MethodCase.VALUESTATEUPDATE => val byteArray = message.getValueStateUpdate.getValue.toByteArray - val valueStateTuple = valueStates(stateName) // The value row is serialized as a byte array, we need to convert it back to a Row - val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3) - valueStateTuple._1.update(valueRow) + val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema, + valueStateInfo.deserializer) + valueStateInfo.valueState.update(valueRow) sendResponse(0) case ValueStateCall.MethodCase.CLEAR => - valueStates(stateName)._1.clear() + valueStateInfo.valueState.clear() + sendResponse(0) + case _ => + throw new IllegalArgumentException("Invalid method call") + } + } + + private[sql] def handleListStateRequest(message: ListStateCall): Unit = { + val stateName = message.getStateName + if (!listStates.contains(stateName)) { + logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.") + sendResponse(1, s"List state $stateName is not initialized.") + return + } + val listStateInfo = listStates(stateName) + val deserializer = if (deserializerForTest != null) { + deserializerForTest + } else { + new TransformWithStateInPandasDeserializer(listStateInfo.deserializer) + } + message.getMethodCase match { + case ListStateCall.MethodCase.EXISTS => + if (listStateInfo.listState.exists()) { + sendResponse(0) + } else { + // Send status code 2 to indicate that the list state doesn't have a value yet. + sendResponse(2, s"state $stateName doesn't exist") + } + case ListStateCall.MethodCase.LISTSTATEPUT => + val rows = deserializer.readArrowBatches(inputStream) + listStateInfo.listState.put(rows.toArray) + sendResponse(0) + case ListStateCall.MethodCase.LISTSTATEGET => + val iteratorId = message.getListStateGet.getIteratorId + var iteratorOption = listStateIterators.get(iteratorId) + if (iteratorOption.isEmpty) { + iteratorOption = Some(listStateInfo.listState.get()) + listStateIterators.put(iteratorId, iteratorOption.get) + } + if (!iteratorOption.get.hasNext) { + sendResponse(2, s"List state $stateName doesn't contain any value.") + return + } else { + sendResponse(0) + } + outputStream.flush() + val arrowStreamWriter = if (arrowStreamWriterForTest != null) { + arrowStreamWriterForTest + } else { + val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId, + errorOnDuplicatedFieldNames, largeVarTypes) + val allocator = ArrowUtils.rootAllocator.newChildAllocator( + s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue) + val root = VectorSchemaRoot.create(arrowSchema, allocator) + new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream), + arrowTransformWithStateInPandasMaxRecordsPerBatch) + } + val listRowSerializer = listStateInfo.serializer + // Only write a single batch in each GET request. Stops writing row if rowCount reaches + // the arrowTransformWithStateInPandasMaxRecordsPerBatch limit. This is to handle a case + // when there are multiple state variables, user tries to access a different state variable + // while the current state variable is not exhausted yet. + var rowCount = 0 + while (iteratorOption.get.hasNext && + rowCount < arrowTransformWithStateInPandasMaxRecordsPerBatch) { + val row = iteratorOption.get.next() + val internalRow = listRowSerializer(row) + arrowStreamWriter.writeRow(internalRow) + rowCount += 1 + } + arrowStreamWriter.finalizeCurrentArrowBatch() + case ListStateCall.MethodCase.APPENDVALUE => + val byteArray = message.getAppendValue.getValue.toByteArray + val newRow = PythonSQLUtils.toJVMRow(byteArray, listStateInfo.schema, + listStateInfo.deserializer) + listStateInfo.listState.appendValue(newRow) + sendResponse(0) + case ListStateCall.MethodCase.APPENDLIST => + val rows = deserializer.readArrowBatches(inputStream) + listStateInfo.listState.appendList(rows.toArray) + sendResponse(0) + case ListStateCall.MethodCase.CLEAR => + listStates(stateName).listState.clear() sendResponse(0) case _ => throw new IllegalArgumentException("Invalid method call") @@ -232,23 +350,54 @@ class TransformWithStateInPandasStateServer( outputStream.write(responseMessageBytes) } - private def initializeValueState( + private def initializeStateVariable( stateName: String, schemaString: String, + stateType: StateVariableType.StateVariableType, ttlDurationMs: Option[Int]): Unit = { - if (!valueStates.contains(stateName)) { - val schema = StructType.fromString(schemaString) - val state = if (ttlDurationMs.isEmpty) { - statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema)) - } else { - statefulProcessorHandle.getValueState( - stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get))) - } - val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer() - valueStates.put(stateName, (state, schema, valueRowDeserializer)) - sendResponse(0) - } else { - sendResponse(1, s"state $stateName already exists") + val schema = StructType.fromString(schemaString) + val expressionEncoder = ExpressionEncoder(schema).resolveAndBind() + stateType match { + case StateVariableType.ValueState => if (!valueStates.contains(stateName)) { + val state = if (ttlDurationMs.isEmpty) { + statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema)) + } else { + statefulProcessorHandle.getValueState( + stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get))) + } + valueStates.put(stateName, + ValueStateInfo(state, schema, expressionEncoder.createDeserializer())) + sendResponse(0) + } else { + sendResponse(1, s"Value state $stateName already exists") + } + case StateVariableType.ListState => if (!listStates.contains(stateName)) { + // TODO(SPARK-49744): Add ttl support for list state. + listStates.put(stateName, + ListStateInfo(statefulProcessorHandle.getListState[Row](stateName, + Encoders.row(schema)), schema, expressionEncoder.createDeserializer(), + expressionEncoder.createSerializer())) + sendResponse(0) + } else { + sendResponse(1, s"List state $stateName already exists") + } } } } + +/** + * Case class to store the information of a value state. + */ +case class ValueStateInfo( + valueState: ValueState[Row], + schema: StructType, + deserializer: ExpressionEncoder.Deserializer[Row]) + +/** + * Case class to store the information of a list state. + */ +case class ListStateInfo( + listState: ListState[Row], + schema: StructType, + deserializer: ExpressionEncoder.Deserializer[Row], + serializer: ExpressionEncoder.Serializer[Row]) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala index 615e1e89f30b8..137e2531f4f46 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala @@ -32,32 +32,59 @@ import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema import org.apache.spark.sql.execution.streaming.{StatefulProcessorHandleImpl, StatefulProcessorHandleState} import org.apache.spark.sql.execution.streaming.state.StateMessage -import org.apache.spark.sql.execution.streaming.state.StateMessage.{Clear, Exists, Get, HandleState, SetHandleState, StateCallCommand, StatefulProcessorCall, ValueStateCall, ValueStateUpdate} -import org.apache.spark.sql.streaming.{TTLConfig, ValueState} +import org.apache.spark.sql.execution.streaming.state.StateMessage.{AppendList, AppendValue, Clear, Exists, Get, HandleState, ListStateCall, ListStateGet, ListStatePut, SetHandleState, StateCallCommand, StatefulProcessorCall, ValueStateCall, ValueStateUpdate} +import org.apache.spark.sql.streaming.{ListState, TTLConfig, ValueState} import org.apache.spark.sql.types.{IntegerType, StructField, StructType} class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with BeforeAndAfterEach { - val valueStateName = "test" - var statefulProcessorHandle: StatefulProcessorHandleImpl = _ + val stateName = "test" + val iteratorId = "testId" + val serverSocket: ServerSocket = mock(classOf[ServerSocket]) + val groupingKeySchema: StructType = StructType(Seq()) + val stateSchema: StructType = StructType(Array(StructField("value", IntegerType))) + // Below byte array is a serialized row with a single integer value 1. + val byteArray: Array[Byte] = Array(0x80.toByte, 0x05.toByte, 0x95.toByte, 0x05.toByte, + 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, + 'K'.toByte, 0x01.toByte, 0x85.toByte, 0x94.toByte, '.'.toByte + ) + + var statefulProcessorHandle: StatefulProcessorHandleImpl = + mock(classOf[StatefulProcessorHandleImpl]) var outputStream: DataOutputStream = _ var valueState: ValueState[Row] = _ + var listState: ListState[Row] = _ var stateServer: TransformWithStateInPandasStateServer = _ - var valueSchema: StructType = _ - var valueDeserializer: ExpressionEncoder.Deserializer[Row] = _ + var stateDeserializer: ExpressionEncoder.Deserializer[Row] = _ + var stateSerializer: ExpressionEncoder.Serializer[Row] = _ + var transformWithStateInPandasDeserializer: TransformWithStateInPandasDeserializer = _ + var arrowStreamWriter: BaseStreamingArrowWriter = _ + var valueStateMap: mutable.HashMap[String, ValueStateInfo] = mutable.HashMap() + var listStateMap: mutable.HashMap[String, ListStateInfo] = mutable.HashMap() override def beforeEach(): Unit = { - val serverSocket = mock(classOf[ServerSocket]) statefulProcessorHandle = mock(classOf[StatefulProcessorHandleImpl]) - val groupingKeySchema = StructType(Seq()) outputStream = mock(classOf[DataOutputStream]) valueState = mock(classOf[ValueState[Row]]) - valueSchema = StructType(Array(StructField("value", IntegerType))) - valueDeserializer = ExpressionEncoder(valueSchema).resolveAndBind().createDeserializer() - val valueStateMap = mutable.HashMap[String, - (ValueState[Row], StructType, ExpressionEncoder.Deserializer[Row])](valueStateName -> - (valueState, valueSchema, valueDeserializer)) + listState = mock(classOf[ListState[Row]]) + stateDeserializer = ExpressionEncoder(stateSchema).resolveAndBind().createDeserializer() + stateSerializer = ExpressionEncoder(stateSchema).resolveAndBind().createSerializer() + valueStateMap = mutable.HashMap[String, ValueStateInfo](stateName -> + ValueStateInfo(valueState, stateSchema, stateDeserializer)) + listStateMap = mutable.HashMap[String, ListStateInfo](stateName -> + ListStateInfo(listState, stateSchema, stateDeserializer, stateSerializer)) + // Iterator map for list state. Please note that `handleImplicitGroupingKeyRequest` would + // reset the iterator map to empty so be careful to call it if you want to access the iterator + // map later. + val listStateIteratorMap = mutable.HashMap[String, Iterator[Row]](iteratorId -> + Iterator(new GenericRowWithSchema(Array(1), stateSchema))) + transformWithStateInPandasDeserializer = mock(classOf[TransformWithStateInPandasDeserializer]) + arrowStreamWriter = mock(classOf[BaseStreamingArrowWriter]) stateServer = new TransformWithStateInPandasStateServer(serverSocket, - statefulProcessorHandle, groupingKeySchema, outputStream, valueStateMap) + statefulProcessorHandle, groupingKeySchema, "", false, false, 2, + outputStream, valueStateMap, transformWithStateInPandasDeserializer, arrowStreamWriter, + listStateMap, listStateIteratorMap) + when(transformWithStateInPandasDeserializer.readArrowBatches(any)) + .thenReturn(Seq(new GenericRowWithSchema(Array(1), stateSchema))) } test("set handle state") { @@ -92,14 +119,14 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo } test("value state exists") { - val message = ValueStateCall.newBuilder().setStateName(valueStateName) + val message = ValueStateCall.newBuilder().setStateName(stateName) .setExists(Exists.newBuilder().build()).build() stateServer.handleValueStateRequest(message) verify(valueState).exists() } test("value state get") { - val message = ValueStateCall.newBuilder().setStateName(valueStateName) + val message = ValueStateCall.newBuilder().setStateName(stateName) .setGet(Get.newBuilder().build()).build() val schema = new StructType().add("value", "int") when(valueState.getOption()).thenReturn(Some(new GenericRowWithSchema(Array(1), schema))) @@ -109,7 +136,7 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo } test("value state get - not exist") { - val message = ValueStateCall.newBuilder().setStateName(valueStateName) + val message = ValueStateCall.newBuilder().setStateName(stateName) .setGet(Get.newBuilder().build()).build() when(valueState.getOption()).thenReturn(None) stateServer.handleValueStateRequest(message) @@ -127,7 +154,7 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo } test("value state clear") { - val message = ValueStateCall.newBuilder().setStateName(valueStateName) + val message = ValueStateCall.newBuilder().setStateName(stateName) .setClear(Clear.newBuilder().build()).build() stateServer.handleValueStateRequest(message) verify(valueState).clear() @@ -135,16 +162,98 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo } test("value state update") { - // Below byte array is a serialized row with a single integer value 1. - val byteArray: Array[Byte] = Array(0x80.toByte, 0x05.toByte, 0x95.toByte, 0x05.toByte, - 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, - 'K'.toByte, 0x01.toByte, 0x85.toByte, 0x94.toByte, '.'.toByte - ) val byteString: ByteString = ByteString.copyFrom(byteArray) - val message = ValueStateCall.newBuilder().setStateName(valueStateName) + val message = ValueStateCall.newBuilder().setStateName(stateName) .setValueStateUpdate(ValueStateUpdate.newBuilder().setValue(byteString).build()).build() stateServer.handleValueStateRequest(message) verify(valueState).update(any[Row]) verify(outputStream).writeInt(0) } + + test("list state exists") { + val message = ListStateCall.newBuilder().setStateName(stateName) + .setExists(Exists.newBuilder().build()).build() + stateServer.handleListStateRequest(message) + verify(listState).exists() + } + + test("list state get - iterator in map") { + val message = ListStateCall.newBuilder().setStateName(stateName) + .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build() + stateServer.handleListStateRequest(message) + verify(listState, times(0)).get() + verify(arrowStreamWriter).writeRow(any) + verify(arrowStreamWriter).finalizeCurrentArrowBatch() + } + + test("list state get - iterator in map with multiple batches") { + val maxRecordsPerBatch = 2 + val message = ListStateCall.newBuilder().setStateName(stateName) + .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build() + val iteratorMap = mutable.HashMap[String, Iterator[Row]](iteratorId -> + Iterator(new GenericRowWithSchema(Array(1), stateSchema), + new GenericRowWithSchema(Array(2), stateSchema), + new GenericRowWithSchema(Array(3), stateSchema), + new GenericRowWithSchema(Array(4), stateSchema))) + stateServer = new TransformWithStateInPandasStateServer(serverSocket, + statefulProcessorHandle, groupingKeySchema, "", false, false, + maxRecordsPerBatch, outputStream, valueStateMap, + transformWithStateInPandasDeserializer, arrowStreamWriter, listStateMap, iteratorMap) + // First call should send 2 records. + stateServer.handleListStateRequest(message) + verify(listState, times(0)).get() + verify(arrowStreamWriter, times(maxRecordsPerBatch)).writeRow(any) + verify(arrowStreamWriter).finalizeCurrentArrowBatch() + // Second call should send the remaining 2 records. + stateServer.handleListStateRequest(message) + verify(listState, times(0)).get() + // Since Mockito's verify counts the total number of calls, the expected number of writeRow call + // should be 2 * maxRecordsPerBatch. + verify(arrowStreamWriter, times(2 * maxRecordsPerBatch)).writeRow(any) + verify(arrowStreamWriter, times(2)).finalizeCurrentArrowBatch() + } + + test("list state get - iterator not in map") { + val maxRecordsPerBatch = 2 + val message = ListStateCall.newBuilder().setStateName(stateName) + .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build() + val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap() + stateServer = new TransformWithStateInPandasStateServer(serverSocket, + statefulProcessorHandle, groupingKeySchema, "", false, false, + maxRecordsPerBatch, outputStream, valueStateMap, + transformWithStateInPandasDeserializer, arrowStreamWriter, listStateMap, iteratorMap) + when(listState.get()).thenReturn(Iterator(new GenericRowWithSchema(Array(1), stateSchema), + new GenericRowWithSchema(Array(2), stateSchema), + new GenericRowWithSchema(Array(3), stateSchema))) + stateServer.handleListStateRequest(message) + verify(listState).get() + // Verify that only maxRecordsPerBatch (2) rows are written to the output stream while still + // having 1 row left in the iterator. + verify(arrowStreamWriter, times(maxRecordsPerBatch)).writeRow(any) + verify(arrowStreamWriter).finalizeCurrentArrowBatch() + } + + test("list state put") { + val message = ListStateCall.newBuilder().setStateName(stateName) + .setListStatePut(ListStatePut.newBuilder().build()).build() + stateServer.handleListStateRequest(message) + verify(transformWithStateInPandasDeserializer).readArrowBatches(any) + verify(listState).put(any) + } + + test("list state append value") { + val byteString: ByteString = ByteString.copyFrom(byteArray) + val message = ListStateCall.newBuilder().setStateName(stateName) + .setAppendValue(AppendValue.newBuilder().setValue(byteString).build()).build() + stateServer.handleListStateRequest(message) + verify(listState).appendValue(any[Row]) + } + + test("list state append list") { + val message = ListStateCall.newBuilder().setStateName(stateName) + .setAppendList(AppendList.newBuilder().build()).build() + stateServer.handleListStateRequest(message) + verify(transformWithStateInPandasDeserializer).readArrowBatches(any) + verify(listState).appendList(any) + } } From a4fb6cbfda228de407e2be83e28c761381576276 Mon Sep 17 00:00:00 2001 From: Nikhil Sheoran <125331115+nikhilsheoran-db@users.noreply.github.com> Date: Wed, 25 Sep 2024 10:29:34 +0800 Subject: [PATCH 127/189] [SPARK-49743][SQL] OptimizeCsvJsonExpr should not change schema fields when pruning GetArrayStructFields ### What changes were proposed in this pull request? - When pruning the schema of the struct in `GetArrayStructFields`, rely on the existing `StructType` to obtain the pruned schema instead of using the accessed field. ### Why are the changes needed? - Fixes a bug in `OptimizeCsvJsonExprs` rule that would have otherwise changed the schema fields of the underlying struct to be extracted. - This would show up as a correctness issue where for a field instead of picking the right values we would have ended up giving null output. ### Does this PR introduce _any_ user-facing change? Yes. The query output would change for the queries of the following type: ``` SELECT from_json('[{"a": '||id||', "b": '|| (2*id) ||'}]', 'array>').a, from_json('[{"a": '||id||', "b": '|| (2*id) ||'}]', 'array>').A FROM range(3) as t ``` Earlier, the result would had been: ``` Array([ArraySeq(0),ArraySeq(null)], [ArraySeq(1),ArraySeq(null)], [ArraySeq(2),ArraySeq(null)]) ``` vs the new result is (verified through spark-shell): ``` Array([ArraySeq(0),ArraySeq(0)], [ArraySeq(1),ArraySeq(1)], [ArraySeq(2),ArraySeq(2)]) ``` ### How was this patch tested? - Added unit tests. - Without this change, the added test would fail as we would have modified the schema from `a` to `A`: ``` - SPARK-49743: prune unnecessary columns from GetArrayStructFields does not change schema *** FAILED *** == FAIL: Plans do not match === !Project [from_json(ArrayType(StructType(StructField(A,IntegerType,true)),true), json#0, Some(America/Los_Angeles)).A AS a#0] Project [from_json(ArrayType(StructType(S tructField(a,IntegerType,true)),true), json#0, Some(America/Los_Angeles)).A AS a#0] +- LocalRelation , [json#0] +- LocalRelation , [json#0] (PlanT est.scala:179) ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48190 from nikhilsheoran-db/SPARK-49743. Authored-by: Nikhil Sheoran <125331115+nikhilsheoran-db@users.noreply.github.com> Signed-off-by: Wenchen Fan --- .../optimizer/OptimizeCsvJsonExprs.scala | 7 ++++--- .../optimizer/OptimizeJsonExprsSuite.scala | 17 +++++++++++++++++ .../org/apache/spark/sql/SQLQuerySuite.scala | 13 +++++++++++++ 3 files changed, 34 insertions(+), 3 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala index 4347137bf68b8..04cc230f99b44 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala @@ -112,9 +112,10 @@ object OptimizeCsvJsonExprs extends Rule[LogicalPlan] { val prunedSchema = StructType(Array(schema(ordinal))) g.copy(child = j.copy(schema = prunedSchema), ordinal = 0) - case g @ GetArrayStructFields(j @ JsonToStructs(schema: ArrayType, _, _, _), _, _, _, _) - if schema.elementType.asInstanceOf[StructType].length > 1 && j.options.isEmpty => - val prunedSchema = ArrayType(StructType(Array(g.field)), g.containsNull) + case g @ GetArrayStructFields(j @ JsonToStructs(ArrayType(schema: StructType, _), + _, _, _), _, ordinal, _, _) if schema.length > 1 && j.options.isEmpty => + // Obtain the pruned schema by picking the `ordinal` field of the struct. + val prunedSchema = ArrayType(StructType(Array(schema(ordinal))), g.containsNull) g.copy(child = j.copy(schema = prunedSchema), ordinal = 0, numFields = 1) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala index c185de4c05d88..eed06da609f8e 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala @@ -307,4 +307,21 @@ class OptimizeJsonExprsSuite extends PlanTest with ExpressionEvalHelper { comparePlans(optimized, query.analyze) } } + + test("SPARK-49743: prune unnecessary columns from GetArrayStructFields does not change schema") { + val options = Map.empty[String, String] + val schema = ArrayType(StructType.fromDDL("a int, b int"), containsNull = true) + + val field = StructField("A", IntegerType) // Instead of "a", use "A" to test case sensitivity. + val query = testRelation2 + .select(GetArrayStructFields( + JsonToStructs(schema, options, $"json"), field, 0, 2, true).as("a")) + val optimized = Optimizer.execute(query.analyze) + + val prunedSchema = ArrayType(StructType.fromDDL("a int"), containsNull = true) + val expected = testRelation2 + .select(GetArrayStructFields( + JsonToStructs(prunedSchema, options, $"json"), field, 0, 1, true).as("a")).analyze + comparePlans(optimized, expected) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index 8176d02dbd02d..e3346684285a9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -4928,6 +4928,19 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark ) ) } + + test("SPARK-49743: OptimizeCsvJsonExpr does not change schema when pruning struct") { + val df = sql(""" + | SELECT + | from_json('[{"a": '||id||', "b": '|| (2*id) ||'}]', 'array>').a, + | from_json('[{"a": '||id||', "b": '|| (2*id) ||'}]', 'array>').A + | FROM + | range(3) as t + |""".stripMargin) + val expectedAnswer = Seq( + Row(Array(0), Array(0)), Row(Array(1), Array(1)), Row(Array(2), Array(2))) + checkAnswer(df, expectedAnswer) + } } case class Foo(bar: Option[String]) From e2d2ab510632cc1948cb6b4500e9da49036a96bd Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Wed, 25 Sep 2024 10:57:44 +0800 Subject: [PATCH 128/189] [SPARK-49552][PYTHON] Add DataFrame API support for new 'randstr' and 'uniform' SQL functions ### What changes were proposed in this pull request? In https://github.com/apache/spark/pull/48004 we added new SQL functions `randstr` and `uniform`. This PR adds DataFrame API support for them. For example, in Scala: ``` sql("create table t(col int not null) using csv") sql("insert into t values (0)") val df = sql("select col from t") df.select(randstr(lit(5), lit(0)).alias("x")).select(length(col("x"))) > 5 df.select(uniform(lit(10), lit(20), lit(0)).alias("x")).selectExpr("x > 5") > true ``` ### Why are the changes needed? This improves DataFrame parity with the SQL API. ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds unit test coverage. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48143 from dtenedor/dataframes-uniform-randstr. Authored-by: Daniel Tenedorio Signed-off-by: Ruifeng Zheng --- .../reference/pyspark.sql/functions.rst | 2 + .../pyspark/sql/connect/functions/builtin.py | 28 +++++ python/pyspark/sql/functions/builtin.py | 92 ++++++++++++++++ python/pyspark/sql/tests/test_functions.py | 21 +++- .../org/apache/spark/sql/functions.scala | 45 ++++++++ .../expressions/randomExpressions.scala | 49 +++++++-- .../spark/sql/DataFrameFunctionsSuite.scala | 104 ++++++++++++++++++ 7 files changed, 331 insertions(+), 10 deletions(-) diff --git a/python/docs/source/reference/pyspark.sql/functions.rst b/python/docs/source/reference/pyspark.sql/functions.rst index 4910a5b59273b..6248e71331656 100644 --- a/python/docs/source/reference/pyspark.sql/functions.rst +++ b/python/docs/source/reference/pyspark.sql/functions.rst @@ -148,6 +148,7 @@ Mathematical Functions try_multiply try_subtract unhex + uniform width_bucket @@ -189,6 +190,7 @@ String Functions overlay position printf + randstr regexp_count regexp_extract regexp_extract_all diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index 6953230f5b42e..27b12fff3c0ac 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -1007,6 +1007,22 @@ def unhex(col: "ColumnOrName") -> Column: unhex.__doc__ = pysparkfuncs.unhex.__doc__ +def uniform( + min: Union[Column, int, float], + max: Union[Column, int, float], + seed: Optional[Union[Column, int]] = None, +) -> Column: + if seed is None: + return _invoke_function_over_columns( + "uniform", lit(min), lit(max), lit(random.randint(0, sys.maxsize)) + ) + else: + return _invoke_function_over_columns("uniform", lit(min), lit(max), lit(seed)) + + +uniform.__doc__ = pysparkfuncs.uniform.__doc__ + + def approxCountDistinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column: warnings.warn("Deprecated in 3.4, use approx_count_distinct instead.", FutureWarning) return approx_count_distinct(col, rsd) @@ -2581,6 +2597,18 @@ def regexp_like(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: regexp_like.__doc__ = pysparkfuncs.regexp_like.__doc__ +def randstr(length: Union[Column, int], seed: Optional[Union[Column, int]] = None) -> Column: + if seed is None: + return _invoke_function_over_columns( + "randstr", lit(length), lit(random.randint(0, sys.maxsize)) + ) + else: + return _invoke_function_over_columns("randstr", lit(length), lit(seed)) + + +randstr.__doc__ = pysparkfuncs.randstr.__doc__ + + def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: return _invoke_function_over_columns("regexp_count", str, regexp) diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 09a286fe7c94e..4ca39562cb20b 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -11973,6 +11973,47 @@ def regexp_like(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: return _invoke_function_over_columns("regexp_like", str, regexp) +@_try_remote_functions +def randstr(length: Union[Column, int], seed: Optional[Union[Column, int]] = None) -> Column: + """Returns a string of the specified length whose characters are chosen uniformly at random from + the following pool of characters: 0-9, a-z, A-Z. The random seed is optional. The string length + must be a constant two-byte or four-byte integer (SMALLINT or INT, respectively). + + .. versionadded:: 4.0.0 + + Parameters + ---------- + length : :class:`~pyspark.sql.Column` or int + Number of characters in the string to generate. + seed : :class:`~pyspark.sql.Column` or int + Optional random number seed to use. + + Returns + ------- + :class:`~pyspark.sql.Column` + The generated random string with the specified length. + + Examples + -------- + >>> spark.createDataFrame([('3',)], ['a']) \\ + ... .select(randstr(lit(5), lit(0)).alias('result')) \\ + ... .selectExpr("length(result) > 0").show() + +--------------------+ + |(length(result) > 0)| + +--------------------+ + | true| + +--------------------+ + """ + length = _enum_to_value(length) + length = lit(length) + if seed is None: + return _invoke_function_over_columns("randstr", length) + else: + seed = _enum_to_value(seed) + seed = lit(seed) + return _invoke_function_over_columns("randstr", length, seed) + + @_try_remote_functions def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column: r"""Returns a count of the number of times that the Java regex pattern `regexp` is matched @@ -12339,6 +12380,57 @@ def unhex(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("unhex", col) +@_try_remote_functions +def uniform( + min: Union[Column, int, float], + max: Union[Column, int, float], + seed: Optional[Union[Column, int]] = None, +) -> Column: + """Returns a random value with independent and identically distributed (i.i.d.) values with the + specified range of numbers. The random seed is optional. The provided numbers specifying the + minimum and maximum values of the range must be constant. If both of these numbers are integers, + then the result will also be an integer. Otherwise if one or both of these are floating-point + numbers, then the result will also be a floating-point number. + + .. versionadded:: 4.0.0 + + Parameters + ---------- + min : :class:`~pyspark.sql.Column`, int, or float + Minimum value in the range. + max : :class:`~pyspark.sql.Column`, int, or float + Maximum value in the range. + seed : :class:`~pyspark.sql.Column` or int + Optional random number seed to use. + + Returns + ------- + :class:`~pyspark.sql.Column` + The generated random number within the specified range. + + Examples + -------- + >>> spark.createDataFrame([('3',)], ['a']) \\ + ... .select(uniform(lit(0), lit(10), lit(0)).alias('result')) \\ + ... .selectExpr("result < 15").show() + +-------------+ + |(result < 15)| + +-------------+ + | true| + +-------------+ + """ + min = _enum_to_value(min) + min = lit(min) + max = _enum_to_value(max) + max = lit(max) + if seed is None: + return _invoke_function_over_columns("uniform", min, max) + else: + seed = _enum_to_value(seed) + seed = lit(seed) + return _invoke_function_over_columns("uniform", min, max, seed) + + @_try_remote_functions def length(col: "ColumnOrName") -> Column: """Computes the character length of string data or number of bytes of binary data. diff --git a/python/pyspark/sql/tests/test_functions.py b/python/pyspark/sql/tests/test_functions.py index a0ab9bc9c7d40..a51156e895c62 100644 --- a/python/pyspark/sql/tests/test_functions.py +++ b/python/pyspark/sql/tests/test_functions.py @@ -29,7 +29,7 @@ from pyspark.sql import Row, Window, functions as F, types from pyspark.sql.avro.functions import from_avro, to_avro from pyspark.sql.column import Column -from pyspark.sql.functions.builtin import nullifzero, zeroifnull +from pyspark.sql.functions.builtin import nullifzero, randstr, uniform, zeroifnull from pyspark.testing.sqlutils import ReusedSQLTestCase, SQLTestUtils from pyspark.testing.utils import have_numpy @@ -1610,6 +1610,25 @@ def test_nullifzero_zeroifnull(self): result = df.select(zeroifnull(df.a).alias("r")).collect() self.assertEqual([Row(r=0), Row(r=1)], result) + def test_randstr_uniform(self): + df = self.spark.createDataFrame([(0,)], ["a"]) + result = df.select(randstr(F.lit(5), F.lit(0)).alias("x")).selectExpr("length(x)").collect() + self.assertEqual([Row(5)], result) + # The random seed is optional. + result = df.select(randstr(F.lit(5)).alias("x")).selectExpr("length(x)").collect() + self.assertEqual([Row(5)], result) + + df = self.spark.createDataFrame([(0,)], ["a"]) + result = ( + df.select(uniform(F.lit(10), F.lit(20), F.lit(0)).alias("x")) + .selectExpr("x > 5") + .collect() + ) + self.assertEqual([Row(True)], result) + # The random seed is optional. + result = df.select(uniform(F.lit(10), F.lit(20)).alias("x")).selectExpr("x > 5").collect() + self.assertEqual([Row(True)], result) + class FunctionsTests(ReusedSQLTestCase, FunctionsTestsMixin): pass diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index ab69789c75f50..93bff22621057 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -1896,6 +1896,26 @@ object functions { */ def randn(): Column = randn(SparkClassUtils.random.nextLong) + /** + * Returns a string of the specified length whose characters are chosen uniformly at random from + * the following pool of characters: 0-9, a-z, A-Z. The string length must be a constant + * two-byte or four-byte integer (SMALLINT or INT, respectively). + * + * @group string_funcs + * @since 4.0.0 + */ + def randstr(length: Column): Column = Column.fn("randstr", length) + + /** + * Returns a string of the specified length whose characters are chosen uniformly at random from + * the following pool of characters: 0-9, a-z, A-Z, with the chosen random seed. The string + * length must be a constant two-byte or four-byte integer (SMALLINT or INT, respectively). + * + * @group string_funcs + * @since 4.0.0 + */ + def randstr(length: Column, seed: Column): Column = Column.fn("randstr", length, seed) + /** * Partition ID. * @@ -3740,6 +3760,31 @@ object functions { */ def stack(cols: Column*): Column = Column.fn("stack", cols: _*) + /** + * Returns a random value with independent and identically distributed (i.i.d.) values with the + * specified range of numbers. The provided numbers specifying the minimum and maximum values of + * the range must be constant. If both of these numbers are integers, then the result will also + * be an integer. Otherwise if one or both of these are floating-point numbers, then the result + * will also be a floating-point number. + * + * @group math_funcs + * @since 4.0.0 + */ + def uniform(min: Column, max: Column): Column = Column.fn("uniform", min, max) + + /** + * Returns a random value with independent and identically distributed (i.i.d.) values with the + * specified range of numbers, with the chosen random seed. The provided numbers specifying the + * minimum and maximum values of the range must be constant. If both of these numbers are + * integers, then the result will also be an integer. Otherwise if one or both of these are + * floating-point numbers, then the result will also be a floating-point number. + * + * @group math_funcs + * @since 4.0.0 + */ + def uniform(min: Column, max: Column, seed: Column): Column = + Column.fn("uniform", min, max, seed) + /** * Returns a random value with independent and identically distributed (i.i.d.) uniformly * distributed values in [0, 1). diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala index f329f8346b0de..ada0a73a67958 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala @@ -206,15 +206,18 @@ object Randn { """, since = "4.0.0", group = "math_funcs") -case class Uniform(min: Expression, max: Expression, seedExpression: Expression) +case class Uniform(min: Expression, max: Expression, seedExpression: Expression, hideSeed: Boolean) extends RuntimeReplaceable with TernaryLike[Expression] with RDG { - def this(min: Expression, max: Expression) = this(min, max, UnresolvedSeed) + def this(min: Expression, max: Expression) = + this(min, max, UnresolvedSeed, hideSeed = true) + def this(min: Expression, max: Expression, seedExpression: Expression) = + this(min, max, seedExpression, hideSeed = false) final override lazy val deterministic: Boolean = false override val nodePatterns: Seq[TreePattern] = Seq(RUNTIME_REPLACEABLE, EXPRESSION_WITH_RANDOM_SEED) - override val dataType: DataType = { + override def dataType: DataType = { val first = min.dataType val second = max.dataType (min.dataType, max.dataType) match { @@ -240,6 +243,10 @@ case class Uniform(min: Expression, max: Expression, seedExpression: Expression) case _ => false } + override def sql: String = { + s"uniform(${min.sql}, ${max.sql}${if (hideSeed) "" else s", ${seedExpression.sql}"})" + } + override def checkInputDataTypes(): TypeCheckResult = { var result: TypeCheckResult = TypeCheckResult.TypeCheckSuccess def requiredType = "integer or floating-point" @@ -277,11 +284,11 @@ case class Uniform(min: Expression, max: Expression, seedExpression: Expression) override def third: Expression = seedExpression override def withNewSeed(newSeed: Long): Expression = - Uniform(min, max, Literal(newSeed, LongType)) + Uniform(min, max, Literal(newSeed, LongType), hideSeed) override def withNewChildrenInternal( newFirst: Expression, newSecond: Expression, newThird: Expression): Expression = - Uniform(newFirst, newSecond, newThird) + Uniform(newFirst, newSecond, newThird, hideSeed) override def replacement: Expression = { if (Seq(min, max, seedExpression).exists(_.dataType == NullType)) { @@ -300,6 +307,13 @@ case class Uniform(min: Expression, max: Expression, seedExpression: Expression) } } +object Uniform { + def apply(min: Expression, max: Expression): Uniform = + Uniform(min, max, UnresolvedSeed, hideSeed = true) + def apply(min: Expression, max: Expression, seedExpression: Expression): Uniform = + Uniform(min, max, seedExpression, hideSeed = false) +} + @ExpressionDescription( usage = """ _FUNC_(length[, seed]) - Returns a string of the specified length whose characters are chosen @@ -315,9 +329,13 @@ case class Uniform(min: Expression, max: Expression, seedExpression: Expression) """, since = "4.0.0", group = "string_funcs") -case class RandStr(length: Expression, override val seedExpression: Expression) +case class RandStr( + length: Expression, override val seedExpression: Expression, hideSeed: Boolean) extends ExpressionWithRandomSeed with BinaryLike[Expression] with Nondeterministic { - def this(length: Expression) = this(length, UnresolvedSeed) + def this(length: Expression) = + this(length, UnresolvedSeed, hideSeed = true) + def this(length: Expression, seedExpression: Expression) = + this(length, seedExpression, hideSeed = false) override def nullable: Boolean = false override def dataType: DataType = StringType @@ -339,9 +357,14 @@ case class RandStr(length: Expression, override val seedExpression: Expression) rng = new XORShiftRandom(seed + partitionIndex) } - override def withNewSeed(newSeed: Long): Expression = RandStr(length, Literal(newSeed, LongType)) + override def withNewSeed(newSeed: Long): Expression = + RandStr(length, Literal(newSeed, LongType), hideSeed) override def withNewChildrenInternal(newFirst: Expression, newSecond: Expression): Expression = - RandStr(newFirst, newSecond) + RandStr(newFirst, newSecond, hideSeed) + + override def sql: String = { + s"randstr(${length.sql}${if (hideSeed) "" else s", ${seedExpression.sql}"})" + } override def checkInputDataTypes(): TypeCheckResult = { var result: TypeCheckResult = TypeCheckResult.TypeCheckSuccess @@ -422,3 +445,11 @@ case class RandStr(length: Expression, override val seedExpression: Expression) isNull = FalseLiteral) } } + +object RandStr { + def apply(length: Expression): RandStr = + RandStr(length, UnresolvedSeed, hideSeed = true) + def apply(length: Expression, seedExpression: Expression): RandStr = + RandStr(length, seedExpression, hideSeed = false) +} + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index 0842b92e5d53c..016803635ff60 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala @@ -411,6 +411,110 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession { checkAnswer(df.select(nvl2(col("b"), col("a"), col("c"))), Seq(Row(null))) } + test("randstr function") { + withTable("t") { + sql("create table t(col int not null) using csv") + sql("insert into t values (0)") + val df = sql("select col from t") + checkAnswer( + df.select(randstr(lit(5), lit(0)).alias("x")).select(length(col("x"))), + Seq(Row(5))) + // The random seed is optional. + checkAnswer( + df.select(randstr(lit(5)).alias("x")).select(length(col("x"))), + Seq(Row(5))) + } + // Here we exercise some error cases. + val df = Seq((0)).toDF("a") + var expr = randstr(lit(10), lit("a")) + checkError( + intercept[AnalysisException](df.select(expr)), + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "\"randstr(10, a)\"", + "paramIndex" -> "second", + "inputSql" -> "\"a\"", + "inputType" -> "\"STRING\"", + "requiredType" -> "INT or SMALLINT"), + context = ExpectedContext( + contextType = QueryContextType.DataFrame, + fragment = "randstr", + objectType = "", + objectName = "", + callSitePattern = "", + startIndex = 0, + stopIndex = 0)) + expr = randstr(col("a"), lit(10)) + checkError( + intercept[AnalysisException](df.select(expr)), + condition = "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + parameters = Map( + "inputName" -> "length", + "inputType" -> "INT or SMALLINT", + "inputExpr" -> "\"a\"", + "sqlExpr" -> "\"randstr(a, 10)\""), + context = ExpectedContext( + contextType = QueryContextType.DataFrame, + fragment = "randstr", + objectType = "", + objectName = "", + callSitePattern = "", + startIndex = 0, + stopIndex = 0)) + } + + test("uniform function") { + withTable("t") { + sql("create table t(col int not null) using csv") + sql("insert into t values (0)") + val df = sql("select col from t") + checkAnswer( + df.select(uniform(lit(10), lit(20), lit(0)).alias("x")).selectExpr("x > 5"), + Seq(Row(true))) + // The random seed is optional. + checkAnswer( + df.select(uniform(lit(10), lit(20)).alias("x")).selectExpr("x > 5"), + Seq(Row(true))) + } + // Here we exercise some error cases. + val df = Seq((0)).toDF("a") + var expr = uniform(lit(10), lit("a")) + checkError( + intercept[AnalysisException](df.select(expr)), + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "\"uniform(10, a)\"", + "paramIndex" -> "second", + "inputSql" -> "\"a\"", + "inputType" -> "\"STRING\"", + "requiredType" -> "integer or floating-point"), + context = ExpectedContext( + contextType = QueryContextType.DataFrame, + fragment = "uniform", + objectType = "", + objectName = "", + callSitePattern = "", + startIndex = 0, + stopIndex = 0)) + expr = uniform(col("a"), lit(10)) + checkError( + intercept[AnalysisException](df.select(expr)), + condition = "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT", + parameters = Map( + "inputName" -> "min", + "inputType" -> "integer or floating-point", + "inputExpr" -> "\"a\"", + "sqlExpr" -> "\"uniform(a, 10)\""), + context = ExpectedContext( + contextType = QueryContextType.DataFrame, + fragment = "uniform", + objectType = "", + objectName = "", + callSitePattern = "", + startIndex = 0, + stopIndex = 0)) + } + test("zeroifnull function") { withTable("t") { // Here we exercise a non-nullable, non-foldable column. From 9aa11d1ee480498de58f0ebd660535effca8fcc6 Mon Sep 17 00:00:00 2001 From: Livia Zhu Date: Wed, 25 Sep 2024 12:42:01 +0900 Subject: [PATCH 129/189] [SPARK-49772][SS] Remove ColumnFamilyOptions and add configs directly to dbOptions in RocksDB ### What changes were proposed in this pull request? To reduce confusion from having vestigial `columnFamilyOptions` value, removed it and added the options directly to `dbOptions`. Also renamed `dbOptions` to `rocksDbOptions` ### Why are the changes needed? Refactoring to simplify and clarify the RocksDB options. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Updated and ensured that existing unit tests pass. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48232 from liviazhu-db/liviazhu-db/rocksdb-options. Lead-authored-by: Livia Zhu Co-authored-by: Jungtaek Lim Signed-off-by: Jungtaek Lim --- .../execution/streaming/state/RocksDB.scala | 39 +++++++++---------- .../streaming/state/RocksDBSuite.scala | 4 +- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala index 4a2aac43b3331..f8d0c8722c3f5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala @@ -115,37 +115,34 @@ class RocksDB( tableFormatConfig.setPinL0FilterAndIndexBlocksInCache(true) } - private[state] val columnFamilyOptions = new ColumnFamilyOptions() + private[state] val rocksDbOptions = new Options() // options to open the RocksDB + + rocksDbOptions.setCreateIfMissing(true) // Set RocksDB options around MemTable memory usage. By default, we let RocksDB // use its internal default values for these settings. if (conf.writeBufferSizeMB > 0L) { - columnFamilyOptions.setWriteBufferSize(conf.writeBufferSizeMB * 1024 * 1024) + rocksDbOptions.setWriteBufferSize(conf.writeBufferSizeMB * 1024 * 1024) } if (conf.maxWriteBufferNumber > 0L) { - columnFamilyOptions.setMaxWriteBufferNumber(conf.maxWriteBufferNumber) + rocksDbOptions.setMaxWriteBufferNumber(conf.maxWriteBufferNumber) } - columnFamilyOptions.setCompressionType(getCompressionType(conf.compression)) - columnFamilyOptions.setMergeOperator(new StringAppendOperator()) - - private val dbOptions = - new Options(new DBOptions(), columnFamilyOptions) // options to open the RocksDB + rocksDbOptions.setCompressionType(getCompressionType(conf.compression)) - dbOptions.setCreateIfMissing(true) - dbOptions.setTableFormatConfig(tableFormatConfig) - dbOptions.setMaxOpenFiles(conf.maxOpenFiles) - dbOptions.setAllowFAllocate(conf.allowFAllocate) - dbOptions.setMergeOperator(new StringAppendOperator()) + rocksDbOptions.setTableFormatConfig(tableFormatConfig) + rocksDbOptions.setMaxOpenFiles(conf.maxOpenFiles) + rocksDbOptions.setAllowFAllocate(conf.allowFAllocate) + rocksDbOptions.setMergeOperator(new StringAppendOperator()) if (conf.boundedMemoryUsage) { - dbOptions.setWriteBufferManager(writeBufferManager) + rocksDbOptions.setWriteBufferManager(writeBufferManager) } private val dbLogger = createLogger() // for forwarding RocksDB native logs to log4j - dbOptions.setStatistics(new Statistics()) - private val nativeStats = dbOptions.statistics() + rocksDbOptions.setStatistics(new Statistics()) + private val nativeStats = rocksDbOptions.statistics() private val workingDir = createTempDir("workingDir") private val fileManager = new RocksDBFileManager(dfsRootDir, createTempDir("fileManager"), @@ -782,7 +779,7 @@ class RocksDB( readOptions.close() writeOptions.close() flushOptions.close() - dbOptions.close() + rocksDbOptions.close() dbLogger.close() synchronized { latestSnapshot.foreach(_.close()) @@ -941,7 +938,7 @@ class RocksDB( private def openDB(): Unit = { assert(db == null) - db = NativeRocksDB.open(dbOptions, workingDir.toString) + db = NativeRocksDB.open(rocksDbOptions, workingDir.toString) logInfo(log"Opened DB with conf ${MDC(LogKeys.CONFIG, conf)}") } @@ -962,7 +959,7 @@ class RocksDB( /** Create a native RocksDB logger that forwards native logs to log4j with correct log levels. */ private def createLogger(): Logger = { - val dbLogger = new Logger(dbOptions.infoLogLevel()) { + val dbLogger = new Logger(rocksDbOptions.infoLogLevel()) { override def log(infoLogLevel: InfoLogLevel, logMsg: String) = { // Map DB log level to log4j levels // Warn is mapped to info because RocksDB warn is too verbose @@ -985,8 +982,8 @@ class RocksDB( dbLogger.setInfoLogLevel(dbLogLevel) // The log level set in dbLogger is effective and the one to dbOptions isn't applied to // customized logger. We still set it as it might show up in RocksDB config file or logging. - dbOptions.setInfoLogLevel(dbLogLevel) - dbOptions.setLogger(dbLogger) + rocksDbOptions.setInfoLogLevel(dbLogLevel) + rocksDbOptions.setLogger(dbLogger) logInfo(log"Set RocksDB native logging level to ${MDC(LogKeys.ROCKS_DB_LOG_LEVEL, dbLogLevel)}") dbLogger } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala index 608a22a284b6c..9fcd2001cce50 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala @@ -526,12 +526,12 @@ class RocksDBSuite extends AlsoTestWithChangelogCheckpointingEnabled with Shared val conf = RocksDBConf().copy(compression = "zstd") withDB(remoteDir, conf = conf, useColumnFamilies = colFamiliesEnabled) { db => - assert(db.columnFamilyOptions.compressionType() == CompressionType.ZSTD_COMPRESSION) + assert(db.rocksDbOptions.compressionType() == CompressionType.ZSTD_COMPRESSION) } // Test the default is LZ4 withDB(remoteDir, conf = RocksDBConf().copy(), useColumnFamilies = colFamiliesEnabled) { db => - assert(db.columnFamilyOptions.compressionType() == CompressionType.LZ4_COMPRESSION) + assert(db.rocksDbOptions.compressionType() == CompressionType.LZ4_COMPRESSION) } } From 5134c68896738179d34e2220ac6171c317900f61 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Wed, 25 Sep 2024 14:29:12 +0900 Subject: [PATCH 130/189] [SPARK-49765][DOCS][PYTHON] Adjust documentation of "spark.sql.pyspark.plotting.max_rows" ### What changes were proposed in this pull request? Adjust documentation of "spark.sql.pyspark.plotting.max_rows". ### Why are the changes needed? Adjust for https://github.com/apache/spark/pull/48218, which eliminates the need for the "spark.sql.pyspark.plotting.sample_ratio" config. ### Does this PR introduce _any_ user-facing change? Doc change only. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48221 from xinrong-meng/conf_doc. Authored-by: Xinrong Meng Signed-off-by: Hyukjin Kwon --- .../main/scala/org/apache/spark/sql/internal/SQLConf.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index c9c227a21cfff..9c46dd8e83ab2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3171,9 +3171,9 @@ object SQLConf { val PYSPARK_PLOT_MAX_ROWS = buildConf("spark.sql.pyspark.plotting.max_rows") - .doc( - "The visual limit on top-n-based plots. If set to 1000, the first 1000 data points " + - "will be used for plotting.") + .doc("The visual limit on plots. If set to 1000 for top-n-based plots (pie, bar, barh), " + + "the first 1000 data points will be used for plotting. For sampled-based plots " + + "(scatter, area, line), 1000 data points will be randomly sampled.") .version("4.0.0") .intConf .createWithDefault(1000) From 46c5accaa55101fe59bce916c17516a70fdfe134 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 25 Sep 2024 19:04:02 +0900 Subject: [PATCH 131/189] [SPARK-49609][PYTHON][TESTS][FOLLOW-UP] Skip Spark Connect tests if dependencies are not found ### What changes were proposed in this pull request? This PR is a followup of https://github.com/apache/spark/pull/48085 that skips the compatibility tests if Spark Connect dependencies are not installed. ### Why are the changes needed? To recover the PyPy3 build https://github.com/apache/spark/actions/runs/11016544408/job/30592416115 which does not have PyArrow installed. ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? Manually. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48239 from HyukjinKwon/SPARK-49609-followup. Authored-by: Hyukjin Kwon Signed-off-by: Haejoon Lee --- python/pyspark/sql/tests/test_connect_compatibility.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/pyspark/sql/tests/test_connect_compatibility.py b/python/pyspark/sql/tests/test_connect_compatibility.py index ca1f828ef4d78..8f3e86f5186a8 100644 --- a/python/pyspark/sql/tests/test_connect_compatibility.py +++ b/python/pyspark/sql/tests/test_connect_compatibility.py @@ -18,6 +18,7 @@ import unittest import inspect +from pyspark.testing.connectutils import should_test_connect, connect_requirement_message from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.classic.dataframe import DataFrame as ClassicDataFrame from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame @@ -172,6 +173,7 @@ def check_missing_methods(classic_cls, connect_cls, cls_name, expected_missing_m ) +@unittest.skipIf(not should_test_connect, connect_requirement_message) class ConnectCompatibilityTests(ConnectCompatibilityTestsMixin, ReusedSQLTestCase): pass From 7f0ecd4221a7043b539fb20a792c00f379a5885e Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Wed, 25 Sep 2024 19:24:05 +0900 Subject: [PATCH 132/189] [SPARK-49764][PYTHON][CONNECT] Support area plots ### What changes were proposed in this pull request? Support area plots with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. Area plots are supported as shown below. ```py >>> from datetime import datetime >>> data = [ ... (3, 5, 20, datetime(2018, 1, 31)), ... (2, 5, 42, datetime(2018, 2, 28)), ... (3, 6, 28, datetime(2018, 3, 31)), ... (9, 12, 62, datetime(2018, 4, 30))] >>> columns = ["sales", "signups", "visits", "date"] >>> df = spark.createDataFrame(data, columns) >>> fig = df.plot.area(x="date", y=["sales", "signups", "visits"]) # df.plot(kind="area", x="date", y=["sales", "signups", "visits"]) >>> fig.show() ``` ![newplot (7)](https://github.com/user-attachments/assets/e603cd99-ce8b-4448-8e1f-cbc093097c45) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48236 from xinrong-meng/plot_area. Authored-by: Xinrong Meng Signed-off-by: Hyukjin Kwon --- python/pyspark/sql/plot/core.py | 35 +++++++++++++++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 35 +++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index 0a3a0101e1898..9f83d00696524 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -93,6 +93,7 @@ def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": class PySparkPlotAccessor: plot_data_map = { + "area": PySparkSampledPlotBase().get_sampled, "bar": PySparkTopNPlotBase().get_top_n, "barh": PySparkTopNPlotBase().get_top_n, "line": PySparkSampledPlotBase().get_sampled, @@ -264,3 +265,37 @@ def scatter(self, x: str, y: str, **kwargs: Any) -> "Figure": >>> df.plot.scatter(x='length', y='width') # doctest: +SKIP """ return self(kind="scatter", x=x, y=y, **kwargs) + + def area(self, x: str, y: str, **kwargs: Any) -> "Figure": + """ + Draw a stacked area plot. + + An area plot displays quantitative data visually. + + Parameters + ---------- + x : str + Name of column to use for the horizontal axis. + y : str or list of str + Name(s) of the column(s) to plot. + **kwargs: Optional + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + >>> from datetime import datetime + >>> data = [ + ... (3, 5, 20, datetime(2018, 1, 31)), + ... (2, 5, 42, datetime(2018, 2, 28)), + ... (3, 6, 28, datetime(2018, 3, 31)), + ... (9, 12, 62, datetime(2018, 4, 30)) + ... ] + >>> columns = ["sales", "signups", "visits", "date"] + >>> df = spark.createDataFrame(data, columns) + >>> df.plot.area(x='date', y=['sales', 'signups', 'visits']) # doctest: +SKIP + """ + return self(kind="area", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py index ccfe1a75424e0..6176525b49550 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -16,6 +16,8 @@ # import unittest +from datetime import datetime + import pyspark.sql.plot # noqa: F401 from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message @@ -34,6 +36,17 @@ def sdf2(self): columns = ["length", "width", "species"] return self.spark.createDataFrame(data, columns) + @property + def sdf3(self): + data = [ + (3, 5, 20, datetime(2018, 1, 31)), + (2, 5, 42, datetime(2018, 2, 28)), + (3, 6, 28, datetime(2018, 3, 31)), + (9, 12, 62, datetime(2018, 4, 30)), + ] + columns = ["sales", "signups", "visits", "date"] + return self.spark.createDataFrame(data, columns) + def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name=""): if kind == "line": self.assertEqual(fig_data["mode"], "lines") @@ -46,6 +59,11 @@ def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name= elif kind == "scatter": self.assertEqual(fig_data["type"], "scatter") self.assertEqual(fig_data["orientation"], "v") + self.assertEqual(fig_data["mode"], "markers") + elif kind == "area": + self.assertEqual(fig_data["type"], "scatter") + self.assertEqual(fig_data["orientation"], "v") + self.assertEqual(fig_data["mode"], "lines") self.assertEqual(fig_data["xaxis"], "x") self.assertEqual(list(fig_data["x"]), expected_x) @@ -98,6 +116,23 @@ def test_scatter_plot(self): "scatter", fig["data"][0], [3.5, 3.0, 3.2, 3.2, 3.0], [5.1, 4.9, 7.0, 6.4, 5.9] ) + def test_area_plot(self): + # single column as vertical axis + fig = self.sdf3.plot(kind="area", x="date", y="sales") + expected_x = [ + datetime(2018, 1, 31, 0, 0), + datetime(2018, 2, 28, 0, 0), + datetime(2018, 3, 31, 0, 0), + datetime(2018, 4, 30, 0, 0), + ] + self._check_fig_data("area", fig["data"][0], expected_x, [3, 2, 3, 9]) + + # multiple columns as vertical axis + fig = self.sdf3.plot.area(x="date", y=["sales", "signups", "visits"]) + self._check_fig_data("area", fig["data"][0], expected_x, [3, 2, 3, 9], "sales") + self._check_fig_data("area", fig["data"][1], expected_x, [5, 5, 6, 12], "signups") + self._check_fig_data("area", fig["data"][2], expected_x, [20, 42, 28, 62], "visits") + class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): pass From e1b2ac55b4b9463824d3f23eb7fbac88ede843d9 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 25 Sep 2024 19:25:47 +0900 Subject: [PATCH 133/189] [SPARK-49767][PS][CONNECT] Refactor the internal function invocation ### What changes were proposed in this pull request? Refactor the internal function invocation ### Why are the changes needed? by introducing a new helper function `_invoke_internal_function_over_columns`, we no longer need to add dedicated internal functions in `PythonSQLUtils` ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48227 from zhengruifeng/py_fn. Authored-by: Ruifeng Zheng Signed-off-by: Hyukjin Kwon --- python/pyspark/pandas/plot/core.py | 2 +- python/pyspark/pandas/spark/functions.py | 175 +++--------------- python/pyspark/pandas/window.py | 3 +- .../spark/sql/api/python/PythonSQLUtils.scala | 43 +---- .../spark/sql/DataFrameSelfJoinSuite.scala | 3 +- .../org/apache/spark/sql/DataFrameSuite.scala | 3 +- 6 files changed, 33 insertions(+), 196 deletions(-) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index 6f036b7669246..7333fae1ad432 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -215,7 +215,7 @@ def compute_hist(psdf, bins): # refers to org.apache.spark.ml.feature.Bucketizer#binarySearchForBuckets def binary_search_for_buckets(value: Column): - index = SF.binary_search(F.lit(bins), value) + index = SF.array_binary_search(F.lit(bins), value) bucket = F.when(index >= 0, index).otherwise(-index - 2) unboundErrMsg = F.lit(f"value %s out of the bins bounds: [{bins[0]}, {bins[-1]}]") return ( diff --git a/python/pyspark/pandas/spark/functions.py b/python/pyspark/pandas/spark/functions.py index 4bcf07f6f6503..4d95466a98e12 100644 --- a/python/pyspark/pandas/spark/functions.py +++ b/python/pyspark/pandas/spark/functions.py @@ -19,197 +19,72 @@ """ from pyspark.sql import Column, functions as F from pyspark.sql.utils import is_remote -from typing import Union +from typing import Union, TYPE_CHECKING +if TYPE_CHECKING: + from pyspark.sql._typing import ColumnOrName -def product(col: Column, dropna: bool) -> Column: + +def _invoke_internal_function_over_columns(name: str, *cols: "ColumnOrName") -> Column: if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit + from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns - return _invoke_function_over_columns( - "pandas_product", - col, - lit(dropna), - ) + return _invoke_function_over_columns(name, *cols) else: + from pyspark.sql.classic.column import _to_seq, _to_java_column from pyspark import SparkContext sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasProduct(col._jc, dropna)) + return Column(sc._jvm.PythonSQLUtils.internalFn(name, _to_seq(sc, cols, _to_java_column))) -def stddev(col: Column, ddof: int) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit - - return _invoke_function_over_columns( - "pandas_stddev", - col, - lit(ddof), - ) +def product(col: Column, dropna: bool) -> Column: + return _invoke_internal_function_over_columns("pandas_product", col, F.lit(dropna)) - else: - from pyspark import SparkContext - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasStddev(col._jc, ddof)) +def stddev(col: Column, ddof: int) -> Column: + return _invoke_internal_function_over_columns("pandas_stddev", col, F.lit(ddof)) def var(col: Column, ddof: int) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit - - return _invoke_function_over_columns( - "pandas_var", - col, - lit(ddof), - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasVariance(col._jc, ddof)) + return _invoke_internal_function_over_columns("pandas_var", col, F.lit(ddof)) def skew(col: Column) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns - - return _invoke_function_over_columns( - "pandas_skew", - col, - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasSkewness(col._jc)) + return _invoke_internal_function_over_columns("pandas_skew", col) def kurt(col: Column) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns - - return _invoke_function_over_columns( - "pandas_kurt", - col, - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasKurtosis(col._jc)) + return _invoke_internal_function_over_columns("pandas_kurt", col) def mode(col: Column, dropna: bool) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit - - return _invoke_function_over_columns( - "pandas_mode", - col, - lit(dropna), - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasMode(col._jc, dropna)) + return _invoke_internal_function_over_columns("pandas_mode", col, F.lit(dropna)) def covar(col1: Column, col2: Column, ddof: int) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit + return _invoke_internal_function_over_columns("pandas_covar", col1, col2, F.lit(ddof)) - return _invoke_function_over_columns( - "pandas_covar", - col1, - col2, - lit(ddof), - ) - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.pandasCovar(col1._jc, col2._jc, ddof)) - - -def ewm(col: Column, alpha: float, ignore_na: bool) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns, lit - - return _invoke_function_over_columns( - "ewm", - col, - lit(alpha), - lit(ignore_na), - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.ewm(col._jc, alpha, ignore_na)) +def ewm(col: Column, alpha: float, ignorena: bool) -> Column: + return _invoke_internal_function_over_columns("ewm", col, F.lit(alpha), F.lit(ignorena)) def null_index(col: Column) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns - - return _invoke_function_over_columns( - "null_index", - col, - ) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.nullIndex(col._jc)) + return _invoke_internal_function_over_columns("null_index", col) def distributed_sequence_id() -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function - - return _invoke_function("distributed_sequence_id") - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.distributed_sequence_id()) + return _invoke_internal_function_over_columns("distributed_sequence_id") def collect_top_k(col: Column, num: int, reverse: bool) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns + return _invoke_internal_function_over_columns("collect_top_k", col, F.lit(num), F.lit(reverse)) - return _invoke_function_over_columns("collect_top_k", col, F.lit(num), F.lit(reverse)) - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.collect_top_k(col._jc, num, reverse)) - - -def binary_search(col: Column, value: Column) -> Column: - if is_remote(): - from pyspark.sql.connect.functions.builtin import _invoke_function_over_columns - - return _invoke_function_over_columns("array_binary_search", col, value) - - else: - from pyspark import SparkContext - - sc = SparkContext._active_spark_context - return Column(sc._jvm.PythonSQLUtils.binary_search(col._jc, value._jc)) +def array_binary_search(col: Column, value: Column) -> Column: + return _invoke_internal_function_over_columns("array_binary_search", col, value) def make_interval(unit: str, e: Union[Column, int, float]) -> Column: diff --git a/python/pyspark/pandas/window.py b/python/pyspark/pandas/window.py index 0aaeb7df89be5..fb5dd29169e91 100644 --- a/python/pyspark/pandas/window.py +++ b/python/pyspark/pandas/window.py @@ -2434,7 +2434,8 @@ def _compute_unified_alpha(self) -> float: if opt_count != 1: raise ValueError("com, span, halflife, and alpha are mutually exclusive") - return unified_alpha + # convert possible numpy.float64 to float for lit function + return float(unified_alpha) @abstractmethod def _apply_as_series_or_frame(self, func: Callable[[Column], Column]) -> FrameLike: diff --git a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala index bc270e6ac64ad..3504f6e76f79d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/api/python/PythonSQLUtils.scala @@ -36,7 +36,6 @@ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.execution.{ExplainMode, QueryExecution} import org.apache.spark.sql.execution.arrow.ArrowConverters import org.apache.spark.sql.execution.python.EvaluatePython -import org.apache.spark.sql.functions.lit import org.apache.spark.sql.internal.ExpressionUtils.{column, expression} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructType} @@ -147,45 +146,6 @@ private[sql] object PythonSQLUtils extends Logging { def castTimestampNTZToLong(c: Column): Column = Column.internalFn("timestamp_ntz_to_long", c) - def ewm(e: Column, alpha: Double, ignoreNA: Boolean): Column = - Column.internalFn("ewm", e, lit(alpha), lit(ignoreNA)) - - def nullIndex(e: Column): Column = Column.internalFn("null_index", e) - - def collect_top_k(e: Column, num: Int, reverse: Boolean): Column = - Column.internalFn("collect_top_k", e, lit(num), lit(reverse)) - - def binary_search(e: Column, value: Column): Column = - Column.internalFn("array_binary_search", e, value) - - def pandasProduct(e: Column, ignoreNA: Boolean): Column = - Column.internalFn("pandas_product", e, lit(ignoreNA)) - - def pandasStddev(e: Column, ddof: Int): Column = - Column.internalFn("pandas_stddev", e, lit(ddof)) - - def pandasVariance(e: Column, ddof: Int): Column = - Column.internalFn("pandas_var", e, lit(ddof)) - - def pandasSkewness(e: Column): Column = - Column.internalFn("pandas_skew", e) - - def pandasKurtosis(e: Column): Column = - Column.internalFn("pandas_kurt", e) - - def pandasMode(e: Column, ignoreNA: Boolean): Column = - Column.internalFn("pandas_mode", e, lit(ignoreNA)) - - def pandasCovar(col1: Column, col2: Column, ddof: Int): Column = - Column.internalFn("pandas_covar", col1, col2, lit(ddof)) - - /** - * A long column that increases one by one. - * This is for 'distributed-sequence' default index in pandas API on Spark. - */ - def distributed_sequence_id(): Column = - Column.internalFn("distributed_sequence_id") - def unresolvedNamedLambdaVariable(name: String): Column = Column(internal.UnresolvedNamedLambdaVariable.apply(name)) @@ -205,6 +165,9 @@ private[sql] object PythonSQLUtils extends Logging { @scala.annotation.varargs def fn(name: String, arguments: Column*): Column = Column.fn(name, arguments: _*) + + @scala.annotation.varargs + def internalFn(name: String, inputs: Column*): Column = Column.internalFn(name, inputs: _*) } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala index 1d7698df2f1be..f0ed2241fd286 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSelfJoinSuite.scala @@ -18,7 +18,6 @@ package org.apache.spark.sql import org.apache.spark.api.python.PythonEvalType -import org.apache.spark.sql.api.python.PythonSQLUtils.distributed_sequence_id import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, PythonUDF, SortOrder} import org.apache.spark.sql.catalyst.plans.logical.{Expand, Generate, ScriptInputOutputSchema, ScriptTransformation, Window => WindowPlan} import org.apache.spark.sql.expressions.Window @@ -405,7 +404,7 @@ class DataFrameSelfJoinSuite extends QueryTest with SharedSparkSession { assertAmbiguousSelfJoin(df12.join(df11, df11("x") === df12("y"))) // Test for AttachDistributedSequence - val df13 = df1.select(distributed_sequence_id().alias("seq"), col("*")) + val df13 = df1.select(Column.internalFn("distributed_sequence_id").alias("seq"), col("*")) val df14 = df13.filter($"value" === "A2") assertAmbiguousSelfJoin(df13.join(df14, df13("key1") === df14("key2"))) assertAmbiguousSelfJoin(df14.join(df13, df13("key1") === df14("key2"))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index e1774cab4a0de..2c0d9e29bb273 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -29,7 +29,6 @@ import org.scalatest.matchers.should.Matchers._ import org.apache.spark.SparkException import org.apache.spark.api.python.PythonEvalType -import org.apache.spark.sql.api.python.PythonSQLUtils.distributed_sequence_id import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} import org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder @@ -2318,7 +2317,7 @@ class DataFrameSuite extends QueryTest test("SPARK-36338: DataFrame.withSequenceColumn should append unique sequence IDs") { val ids = spark.range(10).repartition(5).select( - distributed_sequence_id().alias("default_index"), col("id")) + Column.internalFn("distributed_sequence_id").alias("default_index"), col("id")) assert(ids.collect().map(_.getLong(0)).toSet === Range(0, 10).toSet) assert(ids.take(5).map(_.getLong(0)).toSet === Range(0, 5).toSet) } From c362d500acba5bcf476a2a91ac9b7441ba1e7e2d Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 25 Sep 2024 19:35:33 +0900 Subject: [PATCH 134/189] [SPARK-49775][SQL][TESTS] Make tests of `INVALID_PARAMETER_VALUE.CHARSET` deterministic ### What changes were proposed in this pull request? Make tests of `INVALID_PARAMETER_VALUE.CHARSET` deterministic ### Why are the changes needed? `VALID_CHARSETS` is a Set, so `VALID_CHARSETS.mkString(", ")` is non-deterministic, and cause failures in different testing environments, e.g. ``` org.scalatest.exceptions.TestFailedException: ansi/string-functions.sql Expected "...sets" : "UTF-16LE, U[TF-8, UTF-32, UTF-16BE, UTF-16, US-ASCII, ISO-8859-1]", "functionName...", but got "...sets" : "UTF-16LE, U[S-ASCII, ISO-8859-1, UTF-8, UTF-32, UTF-16BE, UTF-16]", "functionName..." Result did not match for query #93 select encode('hello', 'WINDOWS-1252') at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472) at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471) at org.scalatest.funsuite.AnyFunSuite.newAssertionFailedException(AnyFunSuite.scala:1564) at org.scalatest.Assertions.assertResult(Assertions.scala:847) at org.scalatest.Assertions.assertResult$(Assertions.scala:842) at org.scalatest.funsuite.AnyFunSuite.assertResult(AnyFunSuite.scala:1564) ``` ### Does this PR introduce _any_ user-facing change? No, test only ### How was this patch tested? updated tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48235 from zhengruifeng/sql_test_sort_charset. Authored-by: Ruifeng Zheng Signed-off-by: Hyukjin Kwon --- .../sql/catalyst/util/CharsetProvider.scala | 2 +- .../results/ansi/string-functions.sql.out | 16 ++++++++-------- .../sql-tests/results/string-functions.sql.out | 16 ++++++++-------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala index 0e7fca24e1374..d85673f2ce811 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala @@ -24,7 +24,7 @@ private[sql] object CharsetProvider { final lazy val VALID_CHARSETS = - Set("us-ascii", "iso-8859-1", "utf-8", "utf-16be", "utf-16le", "utf-16", "utf-32") + Array("us-ascii", "iso-8859-1", "utf-8", "utf-16be", "utf-16le", "utf-16", "utf-32").sorted def forName( charset: String, diff --git a/sql/core/src/test/resources/sql-tests/results/ansi/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/ansi/string-functions.sql.out index cf1bce3c0e504..706673606625b 100644 --- a/sql/core/src/test/resources/sql-tests/results/ansi/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/ansi/string-functions.sql.out @@ -842,7 +842,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -860,7 +860,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -878,7 +878,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -896,7 +896,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -1140,7 +1140,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1158,7 +1158,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1208,7 +1208,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1226,7 +1226,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } diff --git a/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out index 14d7b31f8c63f..3f9f24f817f2c 100644 --- a/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out @@ -778,7 +778,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -796,7 +796,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -814,7 +814,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -832,7 +832,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`encode`", "parameter" : "`charset`" } @@ -1076,7 +1076,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1094,7 +1094,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "Windows-xxx", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1144,7 +1144,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } @@ -1162,7 +1162,7 @@ org.apache.spark.SparkIllegalArgumentException "sqlState" : "22023", "messageParameters" : { "charset" : "WINDOWS-1252", - "charsets" : "utf-8, utf-16be, iso-8859-1, utf-16le, utf-16, utf-32, us-ascii", + "charsets" : "iso-8859-1, us-ascii, utf-16, utf-16be, utf-16le, utf-32, utf-8", "functionName" : "`decode`", "parameter" : "`charset`" } From 0ccf53ae6faabc4420317d379da77a299794c84c Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 25 Sep 2024 19:21:36 +0800 Subject: [PATCH 135/189] [SPARK-49609][PYTHON][FOLLOWUP] Correct the typehint for `filter` and `where` ### What changes were proposed in this pull request? Correct the typehint for `filter` and `where` ### Why are the changes needed? the input `str` should not be treated as column name ### Does this PR introduce _any_ user-facing change? doc change ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48244 from zhengruifeng/py_filter_where. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- python/pyspark/sql/classic/dataframe.py | 2 +- python/pyspark/sql/connect/dataframe.py | 2 +- python/pyspark/sql/dataframe.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index 23484fcf0051f..0dd66a9d86545 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -1787,7 +1787,7 @@ def semanticHash(self) -> int: def inputFiles(self) -> List[str]: return list(self._jdf.inputFiles()) - def where(self, condition: "ColumnOrName") -> ParentDataFrame: + def where(self, condition: Union[Column, str]) -> ParentDataFrame: return self.filter(condition) # Two aliases below were added for pandas compatibility many years ago. diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index cb37af8868aad..146cfe11bc502 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -1260,7 +1260,7 @@ def intersectAll(self, other: ParentDataFrame) -> ParentDataFrame: res._cached_schema = self._merge_cached_schema(other) return res - def where(self, condition: "ColumnOrName") -> ParentDataFrame: + def where(self, condition: Union[Column, str]) -> ParentDataFrame: if not isinstance(condition, (str, Column)): raise PySparkTypeError( errorClass="NOT_COLUMN_OR_STR", diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index 2179a844b1e5e..142034583dbd2 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -3351,7 +3351,7 @@ def selectExpr(self, *expr: Union[str, List[str]]) -> "DataFrame": ... @dispatch_df_method - def filter(self, condition: "ColumnOrName") -> "DataFrame": + def filter(self, condition: Union[Column, str]) -> "DataFrame": """Filters rows using the given condition. :func:`where` is an alias for :func:`filter`. @@ -5902,7 +5902,7 @@ def inputFiles(self) -> List[str]: ... @dispatch_df_method - def where(self, condition: "ColumnOrName") -> "DataFrame": + def where(self, condition: Union[Column, str]) -> "DataFrame": """ :func:`where` is an alias for :func:`filter`. From d23023202185f9fd175059caf7499251848c0758 Mon Sep 17 00:00:00 2001 From: Anish Shrigondekar Date: Wed, 25 Sep 2024 22:41:26 +0900 Subject: [PATCH 136/189] [SPARK-49745][SS] Add change to read registered timers through state data source reader ### What changes were proposed in this pull request? Add change to read registered timers through state data source reader ### Why are the changes needed? Without this, users cannot read registered timers per grouping key within the transformWithState operator ### Does this PR introduce _any_ user-facing change? Yes Users can now read registered timers using the following query: ``` val stateReaderDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, ) .option(StateSourceOptions.READ_REGISTERED_TIMERS, true) .load() ``` ### How was this patch tested? Added unit tests ``` [info] Run completed in 20 seconds, 834 milliseconds. [info] Total number of tests run: 4 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 4, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48205 from anishshri-db/task/SPARK-49745. Lead-authored-by: Anish Shrigondekar Co-authored-by: Jungtaek Lim Signed-off-by: Jungtaek Lim --- .../v2/state/StateDataSource.scala | 50 ++++++-- .../v2/state/StatePartitionReader.scala | 5 +- .../v2/state/utils/SchemaUtil.scala | 33 ++++++ .../StateStoreColumnFamilySchemaUtils.scala | 12 ++ .../streaming/StateTypesEncoderUtils.scala | 3 + .../StatefulProcessorHandleImpl.scala | 16 +++ .../execution/streaming/TimerStateImpl.scala | 9 ++ .../TransformWithStateVariableUtils.scala | 6 +- .../v2/state/StateDataSourceReadSuite.scala | 19 +++ ...ateDataSourceTransformWithStateSuite.scala | 109 +++++++++++++++++- .../TransformWithValueStateTTLSuite.scala | 21 +++- 11 files changed, 263 insertions(+), 20 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala index 429464ea5438d..39bc4dd9fb9c8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala @@ -29,15 +29,16 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.DataSourceOptions import org.apache.spark.sql.connector.catalog.{Table, TableProvider} import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.{JoinSideValues, STATE_VAR_NAME} +import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.{JoinSideValues, READ_REGISTERED_TIMERS, STATE_VAR_NAME} import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.JoinSideValues.JoinSideValues import org.apache.spark.sql.execution.datasources.v2.state.metadata.{StateMetadataPartitionReader, StateMetadataTableEntry} import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil -import org.apache.spark.sql.execution.streaming.{CommitLog, OffsetSeqLog, OffsetSeqMetadata, TransformWithStateOperatorProperties, TransformWithStateVariableInfo} +import org.apache.spark.sql.execution.streaming.{CommitLog, OffsetSeqLog, OffsetSeqMetadata, TimerStateUtils, TransformWithStateOperatorProperties, TransformWithStateVariableInfo} import org.apache.spark.sql.execution.streaming.StreamingCheckpointConstants.{DIR_NAME_COMMITS, DIR_NAME_OFFSETS, DIR_NAME_STATE} import org.apache.spark.sql.execution.streaming.StreamingSymmetricHashJoinHelper.{LeftSide, RightSide} import org.apache.spark.sql.execution.streaming.state.{KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateSchemaCompatibilityChecker, StateStore, StateStoreColFamilySchema, StateStoreConf, StateStoreId, StateStoreProviderId} import org.apache.spark.sql.sources.DataSourceRegister +import org.apache.spark.sql.streaming.TimeMode import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.SerializableConfiguration @@ -132,7 +133,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging sourceOptions: StateSourceOptions, stateStoreMetadata: Array[StateMetadataTableEntry]): Unit = { val twsShortName = "transformWithStateExec" - if (sourceOptions.stateVarName.isDefined) { + if (sourceOptions.stateVarName.isDefined || sourceOptions.readRegisteredTimers) { // Perform checks for transformWithState operator in case state variable name is provided require(stateStoreMetadata.size == 1) val opMetadata = stateStoreMetadata.head @@ -153,10 +154,21 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging "No state variable names are defined for the transformWithState operator") } + val twsOperatorProperties = TransformWithStateOperatorProperties.fromJson(operatorProperties) + val timeMode = twsOperatorProperties.timeMode + if (sourceOptions.readRegisteredTimers && timeMode == TimeMode.None().toString) { + throw StateDataSourceErrors.invalidOptionValue(READ_REGISTERED_TIMERS, + "Registered timers are not available in TimeMode=None.") + } + // if the state variable is not one of the defined/available state variables, then we // fail the query - val stateVarName = sourceOptions.stateVarName.get - val twsOperatorProperties = TransformWithStateOperatorProperties.fromJson(operatorProperties) + val stateVarName = if (sourceOptions.readRegisteredTimers) { + TimerStateUtils.getTimerStateVarName(timeMode) + } else { + sourceOptions.stateVarName.get + } + val stateVars = twsOperatorProperties.stateVariables if (stateVars.filter(stateVar => stateVar.stateName == stateVarName).size != 1) { throw StateDataSourceErrors.invalidOptionValue(STATE_VAR_NAME, @@ -196,9 +208,10 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging var keyStateEncoderSpecOpt: Option[KeyStateEncoderSpec] = None var stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema] = None var transformWithStateVariableInfoOpt: Option[TransformWithStateVariableInfo] = None + var timeMode: String = TimeMode.None.toString if (sourceOptions.joinSide == JoinSideValues.none) { - val stateVarName = sourceOptions.stateVarName + var stateVarName = sourceOptions.stateVarName .getOrElse(StateStore.DEFAULT_COL_FAMILY_NAME) // Read the schema file path from operator metadata version v2 onwards @@ -208,6 +221,12 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging val storeMetadataEntry = storeMetadata.head val operatorProperties = TransformWithStateOperatorProperties.fromJson( storeMetadataEntry.operatorPropertiesJson) + timeMode = operatorProperties.timeMode + + if (sourceOptions.readRegisteredTimers) { + stateVarName = TimerStateUtils.getTimerStateVarName(timeMode) + } + val stateVarInfoList = operatorProperties.stateVariables .filter(stateVar => stateVar.stateName == stateVarName) require(stateVarInfoList.size == 1, s"Failed to find unique state variable info " + @@ -304,6 +323,7 @@ case class StateSourceOptions( fromSnapshotOptions: Option[FromSnapshotOptions], readChangeFeedOptions: Option[ReadChangeFeedOptions], stateVarName: Option[String], + readRegisteredTimers: Boolean, flattenCollectionTypes: Boolean) { def stateCheckpointLocation: Path = new Path(resolvedCpLocation, DIR_NAME_STATE) @@ -336,6 +356,7 @@ object StateSourceOptions extends DataSourceOptions { val CHANGE_START_BATCH_ID = newOption("changeStartBatchId") val CHANGE_END_BATCH_ID = newOption("changeEndBatchId") val STATE_VAR_NAME = newOption("stateVarName") + val READ_REGISTERED_TIMERS = newOption("readRegisteredTimers") val FLATTEN_COLLECTION_TYPES = newOption("flattenCollectionTypes") object JoinSideValues extends Enumeration { @@ -377,6 +398,19 @@ object StateSourceOptions extends DataSourceOptions { val stateVarName = Option(options.get(STATE_VAR_NAME)) .map(_.trim) + val readRegisteredTimers = try { + Option(options.get(READ_REGISTERED_TIMERS)) + .map(_.toBoolean).getOrElse(false) + } catch { + case _: IllegalArgumentException => + throw StateDataSourceErrors.invalidOptionValue(READ_REGISTERED_TIMERS, + "Boolean value is expected") + } + + if (readRegisteredTimers && stateVarName.isDefined) { + throw StateDataSourceErrors.conflictOptions(Seq(READ_REGISTERED_TIMERS, STATE_VAR_NAME)) + } + val flattenCollectionTypes = try { Option(options.get(FLATTEN_COLLECTION_TYPES)) .map(_.toBoolean).getOrElse(true) @@ -489,8 +523,8 @@ object StateSourceOptions extends DataSourceOptions { StateSourceOptions( resolvedCpLocation, batchId.get, operatorId, storeName, joinSide, - readChangeFeed, fromSnapshotOptions, readChangeFeedOptions, stateVarName, - flattenCollectionTypes) + readChangeFeed, fromSnapshotOptions, readChangeFeedOptions, + stateVarName, readRegisteredTimers, flattenCollectionTypes) } private def resolvedCheckpointLocation( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala index ae12b18c1f627..d77d97f0057fb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala @@ -107,6 +107,8 @@ abstract class StatePartitionReaderBase( useColumnFamilies = useColFamilies, storeConf, hadoopConf.value, useMultipleValuesPerKey = useMultipleValuesPerKey) + val isInternal = partition.sourceOptions.readRegisteredTimers + if (useColFamilies) { val store = provider.getStore(partition.sourceOptions.batchId + 1) require(stateStoreColFamilySchemaOpt.isDefined) @@ -117,7 +119,8 @@ abstract class StatePartitionReaderBase( stateStoreColFamilySchema.keySchema, stateStoreColFamilySchema.valueSchema, stateStoreColFamilySchema.keyStateEncoderSpec.get, - useMultipleValuesPerKey = useMultipleValuesPerKey) + useMultipleValuesPerKey = useMultipleValuesPerKey, + isInternal = isInternal) } provider } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala index dc0d6af951143..c337d548fa42b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala @@ -230,6 +230,7 @@ object SchemaUtil { "map_value" -> classOf[MapType], "user_map_key" -> classOf[StructType], "user_map_value" -> classOf[StructType], + "expiration_timestamp_ms" -> classOf[LongType], "partition_id" -> classOf[IntegerType]) val expectedFieldNames = if (sourceOptions.readChangeFeed) { @@ -256,6 +257,9 @@ object SchemaUtil { Seq("key", "map_value", "partition_id") } + case TimerState => + Seq("key", "expiration_timestamp_ms", "partition_id") + case _ => throw StateDataSourceErrors .internalError(s"Unsupported state variable type $stateVarType") @@ -322,6 +326,14 @@ object SchemaUtil { .add("partition_id", IntegerType) } + case TimerState => + val groupingKeySchema = SchemaUtil.getSchemaAsDataType( + stateStoreColFamilySchema.keySchema, "key") + new StructType() + .add("key", groupingKeySchema) + .add("expiration_timestamp_ms", LongType) + .add("partition_id", IntegerType) + case _ => throw StateDataSourceErrors.internalError(s"Unsupported state variable type $stateVarType") } @@ -407,9 +419,30 @@ object SchemaUtil { unifyMapStateRowPair(store.iterator(stateVarName), compositeKeySchema, partitionId, stateSourceOptions) + case StateVariableType.TimerState => + store + .iterator(stateVarName) + .map { pair => + unifyTimerRow(pair.key, compositeKeySchema, partitionId) + } + case _ => throw new IllegalStateException( s"Unsupported state variable type: $stateVarType") } } + + private def unifyTimerRow( + rowKey: UnsafeRow, + groupingKeySchema: StructType, + partitionId: Int): InternalRow = { + val groupingKey = rowKey.get(0, groupingKeySchema).asInstanceOf[UnsafeRow] + val expirationTimestamp = rowKey.getLong(1) + + val row = new GenericInternalRow(3) + row.update(0, groupingKey) + row.update(1, expirationTimestamp) + row.update(2, partitionId) + row + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala index 99229c6132eb2..7da8408f98b0f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala @@ -20,6 +20,7 @@ import org.apache.spark.sql.Encoder import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchemaUtils._ import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateStoreColFamilySchema} +import org.apache.spark.sql.types.StructType object StateStoreColumnFamilySchemaUtils { @@ -61,4 +62,15 @@ object StateStoreColumnFamilySchemaUtils { Some(PrefixKeyScanStateEncoderSpec(compositeKeySchema, 1)), Some(userKeyEnc.schema)) } + + def getTimerStateSchema( + stateName: String, + keySchema: StructType, + valSchema: StructType): StateStoreColFamilySchema = { + StateStoreColFamilySchema( + stateName, + keySchema, + valSchema, + Some(PrefixKeyScanStateEncoderSpec(keySchema, 1))) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala index 1f5ad2fc85470..b70f9699195d4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala @@ -288,6 +288,9 @@ class TimerKeyEncoder(keyExprEnc: ExpressionEncoder[Any]) { .add("key", new StructType(keyExprEnc.schema.fields)) .add("expiryTimestampMs", LongType, nullable = false) + val schemaForValueRow: StructType = + StructType(Array(StructField("__dummy__", NullType))) + private val keySerializer = keyExprEnc.createSerializer() private val keyDeserializer = keyExprEnc.resolveAndBind().createDeserializer() private val prefixKeyProjection = UnsafeProjection.create(schemaForPrefixKey) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala index 942d395dec0e2..8beacbec7e6ef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala @@ -308,6 +308,12 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi private val stateVariableInfos: mutable.Map[String, TransformWithStateVariableInfo] = new mutable.HashMap[String, TransformWithStateVariableInfo]() + // If timeMode is not None, add a timer column family schema to the operator metadata so that + // registered timers can be read using the state data source reader. + if (timeMode != TimeMode.None()) { + addTimerColFamily() + } + def getColumnFamilySchemas: Map[String, StateStoreColFamilySchema] = columnFamilySchemas.toMap def getStateVariableInfos: Map[String, TransformWithStateVariableInfo] = stateVariableInfos.toMap @@ -318,6 +324,16 @@ class DriverStatefulProcessorHandleImpl(timeMode: TimeMode, keyExprEnc: Expressi } } + private def addTimerColFamily(): Unit = { + val stateName = TimerStateUtils.getTimerStateVarName(timeMode.toString) + val timerEncoder = new TimerKeyEncoder(keyExprEnc) + val colFamilySchema = StateStoreColumnFamilySchemaUtils. + getTimerStateSchema(stateName, timerEncoder.schemaForKeyRow, timerEncoder.schemaForValueRow) + columnFamilySchemas.put(stateName, colFamilySchema) + val stateVariableInfo = TransformWithStateVariableUtils.getTimerState(stateName) + stateVariableInfos.put(stateName, stateVariableInfo) + } + override def getValueState[T](stateName: String, valEncoder: Encoder[T]): ValueState[T] = { verifyStateVarOperations("get_value_state", PRE_INIT) val colFamilySchema = StateStoreColumnFamilySchemaUtils. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala index 82a4226fcfd54..d0fbaf6600609 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TimerStateImpl.scala @@ -34,6 +34,15 @@ object TimerStateUtils { val EVENT_TIMERS_STATE_NAME = "$eventTimers" val KEY_TO_TIMESTAMP_CF = "_keyToTimestamp" val TIMESTAMP_TO_KEY_CF = "_timestampToKey" + + def getTimerStateVarName(timeMode: String): String = { + assert(timeMode == TimeMode.EventTime.toString || timeMode == TimeMode.ProcessingTime.toString) + if (timeMode == TimeMode.EventTime.toString) { + TimerStateUtils.EVENT_TIMERS_STATE_NAME + TimerStateUtils.KEY_TO_TIMESTAMP_CF + } else { + TimerStateUtils.PROC_TIMERS_STATE_NAME + TimerStateUtils.KEY_TO_TIMESTAMP_CF + } + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala index 0a32564f973a3..4a192b3e51c71 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala @@ -43,12 +43,16 @@ object TransformWithStateVariableUtils { def getMapState(stateName: String, ttlEnabled: Boolean): TransformWithStateVariableInfo = { TransformWithStateVariableInfo(stateName, StateVariableType.MapState, ttlEnabled) } + + def getTimerState(stateName: String): TransformWithStateVariableInfo = { + TransformWithStateVariableInfo(stateName, StateVariableType.TimerState, ttlEnabled = false) + } } // Enum of possible State Variable types object StateVariableType extends Enumeration { type StateVariableType = Value - val ValueState, ListState, MapState = Value + val ValueState, ListState, MapState, TimerState = Value } case class TransformWithStateVariableInfo( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala index 8707facc4c126..5f55848d540df 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceReadSuite.scala @@ -288,6 +288,25 @@ class StateDataSourceNegativeTestSuite extends StateDataSourceTestBase { } } + test("ERROR: trying to specify state variable name along with " + + "readRegisteredTimers should fail") { + withTempDir { tempDir => + val exc = intercept[StateDataSourceConflictOptions] { + spark.read.format("statestore") + // trick to bypass getting the last committed batch before validating operator ID + .option(StateSourceOptions.BATCH_ID, 0) + .option(StateSourceOptions.STATE_VAR_NAME, "test") + .option(StateSourceOptions.READ_REGISTERED_TIMERS, true) + .load(tempDir.getAbsolutePath) + } + checkError(exc, "STDS_CONFLICT_OPTIONS", "42613", + Map("options" -> + s"['${ + StateSourceOptions.READ_REGISTERED_TIMERS + }', '${StateSourceOptions.STATE_VAR_NAME}']")) + } + } + test("ERROR: trying to specify non boolean value for " + "flattenCollectionTypes") { withTempDir { tempDir => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala index 69df86fd5f746..bd047d1132fbe 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala @@ -21,9 +21,9 @@ import java.time.Duration import org.apache.spark.sql.{Encoders, Row} import org.apache.spark.sql.execution.streaming.MemoryStream import org.apache.spark.sql.execution.streaming.state.{AlsoTestWithChangelogCheckpointingEnabled, RocksDBStateStoreProvider, TestClass} -import org.apache.spark.sql.functions.explode +import org.apache.spark.sql.functions.{explode, timestamp_seconds} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.streaming.{ExpiredTimerInfo, InputMapRow, ListState, MapInputEvent, MapOutputEvent, MapStateTTLProcessor, OutputMode, RunningCountStatefulProcessor, StatefulProcessor, StateStoreMetricsTest, TestMapStateProcessor, TimeMode, TimerValues, TransformWithStateSuiteUtils, Trigger, TTLConfig, ValueState} +import org.apache.spark.sql.streaming.{ExpiredTimerInfo, InputMapRow, ListState, MapInputEvent, MapOutputEvent, MapStateTTLProcessor, MaxEventTimeStatefulProcessor, OutputMode, RunningCountStatefulProcessor, RunningCountStatefulProcessorWithProcTimeTimerUpdates, StatefulProcessor, StateStoreMetricsTest, TestMapStateProcessor, TimeMode, TimerValues, TransformWithStateSuiteUtils, Trigger, TTLConfig, ValueState} import org.apache.spark.sql.streaming.util.StreamManualClock /** Stateful processor of single value state var with non-primitive type */ @@ -176,8 +176,19 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest assert(ex.isInstanceOf[StateDataSourceInvalidOptionValue]) assert(ex.getMessage.contains("State variable non-exist is not defined")) - // TODO: this should be removed when readChangeFeed is supported for value state + // Verify that trying to read timers in TimeMode as None fails val ex1 = intercept[Exception] { + spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.READ_REGISTERED_TIMERS, true) + .load() + } + assert(ex1.isInstanceOf[StateDataSourceInvalidOptionValue]) + assert(ex1.getMessage.contains("Registered timers are not available")) + + // TODO: this should be removed when readChangeFeed is supported for value state + val ex2 = intercept[Exception] { spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) @@ -186,7 +197,7 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest .option(StateSourceOptions.CHANGE_START_BATCH_ID, 0) .load() } - assert(ex1.isInstanceOf[StateDataSourceConflictOptions]) + assert(ex2.isInstanceOf[StateDataSourceConflictOptions]) } } } @@ -563,4 +574,94 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest } } } + + test("state data source - processing-time timers integration") { + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString) { + withTempDir { tempDir => + val clock = new StreamManualClock + + val inputData = MemoryStream[String] + val result = inputData.toDS() + .groupByKey(x => x) + .transformWithState( + new RunningCountStatefulProcessorWithProcTimeTimerUpdates(), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update())( + StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock, + checkpointLocation = tempDir.getCanonicalPath), + AddData(inputData, "a"), + AdvanceManualClock(1 * 1000), + CheckNewAnswer(("a", "1")), // at batch 0, ts = 1, timer = "a" -> [6] (= 1 + 5) + AddData(inputData, "a"), + AdvanceManualClock(2 * 1000), + CheckNewAnswer(("a", "2")), // at batch 1, ts = 3, timer = "a" -> [10.5] (3 + 7.5) + StopStream) + + val stateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.READ_REGISTERED_TIMERS, true) + .load() + + val resultDf = stateReaderDf.selectExpr( + "key.value AS groupingKey", + "expiration_timestamp_ms AS expiryTimestamp", + "partition_id") + + checkAnswer(resultDf, + Seq(Row("a", 10500L, 0))) + } + } + } + + test("state data source - event-time timers integration") { + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString) { + withTempDir { tempDir => + val inputData = MemoryStream[(String, Int)] + val result = + inputData.toDS() + .select($"_1".as("key"), timestamp_seconds($"_2").as("eventTime")) + .withWatermark("eventTime", "10 seconds") + .as[(String, Long)] + .groupByKey(_._1) + .transformWithState( + new MaxEventTimeStatefulProcessor(), + TimeMode.EventTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update())( + StartStream(checkpointLocation = tempDir.getCanonicalPath), + + AddData(inputData, ("a", 11), ("a", 13), ("a", 15)), + // Max event time = 15. Timeout timestamp for "a" = 15 + 5 = 20. Watermark = 15 - 10 = 5. + CheckNewAnswer(("a", 15)), // Output = max event time of a + + AddData(inputData, ("a", 4)), // Add data older than watermark for "a" + CheckNewAnswer(), // No output as data should get filtered by watermark + StopStream) + + val stateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.READ_REGISTERED_TIMERS, true) + .load() + + val resultDf = stateReaderDf.selectExpr( + "key.value AS groupingKey", + "expiration_timestamp_ms AS expiryTimestamp", + "partition_id") + + checkAnswer(resultDf, + Seq(Row("a", 20000L, 0))) + } + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala index 45056d104e84e..1fbeaeb817bd9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala @@ -23,7 +23,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.internal.Logging import org.apache.spark.sql.Encoders -import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, ListStateImplWithTTL, MapStateImplWithTTL, MemoryStream, ValueStateImpl, ValueStateImplWithTTL} +import org.apache.spark.sql.execution.streaming.{CheckpointFileManager, ListStateImplWithTTL, MapStateImplWithTTL, MemoryStream, TimerStateUtils, ValueStateImpl, ValueStateImplWithTTL} import org.apache.spark.sql.execution.streaming.state._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.util.StreamManualClock @@ -265,7 +265,16 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest { val fm = CheckpointFileManager.create(stateSchemaPath, hadoopConf) val keySchema = new StructType().add("value", StringType) + val schemaForKeyRow: StructType = new StructType() + .add("key", new StructType(keySchema.fields)) + .add("expiryTimestampMs", LongType, nullable = false) + val schemaForValueRow: StructType = StructType(Array(StructField("__dummy__", NullType))) val schema0 = StateStoreColFamilySchema( + TimerStateUtils.getTimerStateVarName(TimeMode.ProcessingTime().toString), + schemaForKeyRow, + schemaForValueRow, + Some(PrefixKeyScanStateEncoderSpec(schemaForKeyRow, 1))) + val schema1 = StateStoreColFamilySchema( "valueStateTTL", keySchema, new StructType().add("value", @@ -275,14 +284,14 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest { Some(NoPrefixKeyStateEncoderSpec(keySchema)), None ) - val schema1 = StateStoreColFamilySchema( + val schema2 = StateStoreColFamilySchema( "valueState", keySchema, new StructType().add("value", IntegerType, false), Some(NoPrefixKeyStateEncoderSpec(keySchema)), None ) - val schema2 = StateStoreColFamilySchema( + val schema3 = StateStoreColFamilySchema( "listState", keySchema, new StructType().add("value", @@ -300,7 +309,7 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest { val compositeKeySchema = new StructType() .add("key", new StructType().add("value", StringType)) .add("userKey", userKeySchema) - val schema3 = StateStoreColFamilySchema( + val schema4 = StateStoreColFamilySchema( "mapState", compositeKeySchema, new StructType().add("value", @@ -351,9 +360,9 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest { q.lastProgress.stateOperators.head.customMetrics .get("numMapStateWithTTLVars").toInt) - assert(colFamilySeq.length == 4) + assert(colFamilySeq.length == 5) assert(colFamilySeq.map(_.toString).toSet == Set( - schema0, schema1, schema2, schema3 + schema0, schema1, schema2, schema3, schema4 ).map(_.toString)) }, StopStream From 983f6f434af335b9270a0748dc5b4b18c7dc4846 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Wed, 25 Sep 2024 07:50:20 -0700 Subject: [PATCH 137/189] [SPARK-49746][BUILD] Upgrade Scala to 2.13.15 ### What changes were proposed in this pull request? The pr aims to upgrade `scala` from `2.13.14` to `2.13.15`. ### Why are the changes needed? https://contributors.scala-lang.org/t/scala-2-13-15-release-planning/6649 image **Note: since 2.13.15, "-Wconf:cat=deprecation:wv,any:e" no longer takes effect and needs to be changed to "-Wconf:any:e", "-Wconf:cat=deprecation:wv", please refer to the details: https://github.com/scala/scala/pull/10708** ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48192 from panbingkun/SPARK-49746. Lead-authored-by: panbingkun Co-authored-by: YangJie Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 8 ++++---- docs/_config.yml | 2 +- pom.xml | 7 ++++--- project/SparkBuild.scala | 6 +++++- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 88526995293f5..19b8a237d30aa 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -144,7 +144,7 @@ jetty-util-ajax/11.0.23//jetty-util-ajax-11.0.23.jar jetty-util/11.0.23//jetty-util-11.0.23.jar jjwt-api/0.12.6//jjwt-api-0.12.6.jar jline/2.14.6//jline-2.14.6.jar -jline/3.25.1//jline-3.25.1.jar +jline/3.26.3//jline-3.26.3.jar jna/5.14.0//jna-5.14.0.jar joda-time/2.13.0//joda-time-2.13.0.jar jodd-core/3.5.2//jodd-core-3.5.2.jar @@ -252,11 +252,11 @@ py4j/0.10.9.7//py4j-0.10.9.7.jar remotetea-oncrpc/1.1.2//remotetea-oncrpc-1.1.2.jar rocksdbjni/9.5.2//rocksdbjni-9.5.2.jar scala-collection-compat_2.13/2.7.0//scala-collection-compat_2.13-2.7.0.jar -scala-compiler/2.13.14//scala-compiler-2.13.14.jar -scala-library/2.13.14//scala-library-2.13.14.jar +scala-compiler/2.13.15//scala-compiler-2.13.15.jar +scala-library/2.13.15//scala-library-2.13.15.jar scala-parallel-collections_2.13/1.0.4//scala-parallel-collections_2.13-1.0.4.jar scala-parser-combinators_2.13/2.4.0//scala-parser-combinators_2.13-2.4.0.jar -scala-reflect/2.13.14//scala-reflect-2.13.14.jar +scala-reflect/2.13.15//scala-reflect-2.13.15.jar scala-xml_2.13/2.3.0//scala-xml_2.13-2.3.0.jar slf4j-api/2.0.16//slf4j-api-2.0.16.jar snakeyaml-engine/2.7//snakeyaml-engine-2.7.jar diff --git a/docs/_config.yml b/docs/_config.yml index e74eda0470417..089d6bf2097b8 100644 --- a/docs/_config.yml +++ b/docs/_config.yml @@ -22,7 +22,7 @@ include: SPARK_VERSION: 4.0.0-SNAPSHOT SPARK_VERSION_SHORT: 4.0.0 SCALA_BINARY_VERSION: "2.13" -SCALA_VERSION: "2.13.14" +SCALA_VERSION: "2.13.15" SPARK_ISSUE_TRACKER_URL: https://issues.apache.org/jira/browse/SPARK SPARK_GITHUB_URL: https://github.com/apache/spark # Before a new release, we should: diff --git a/pom.xml b/pom.xml index 131e754da8157..f3dc92426ac4e 100644 --- a/pom.xml +++ b/pom.xml @@ -169,7 +169,7 @@ 3.2.2 4.4 - 2.13.14 + 2.13.15 2.13 2.2.0 4.9.1 @@ -226,7 +226,7 @@ and ./python/packaging/connect/setup.py too. --> 17.0.0 - 3.0.0-M2 + 3.0.0 0.12.6 @@ -3051,7 +3051,8 @@ -explaintypes -release 17 - -Wconf:cat=deprecation:wv,any:e + -Wconf:any:e + -Wconf:cat=deprecation:wv -Wunused:imports -Wconf:cat=scaladoc:wv -Wconf:msg=^(?=.*?method|value|type|object|trait|inheritance)(?=.*?deprecated)(?=.*?since 2.13).+$:e diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index 2f390cb70baa8..82950fb30287a 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -234,7 +234,11 @@ object SparkBuild extends PomBuild { // replace -Xfatal-warnings with fine-grained configuration, since 2.13.2 // verbose warning on deprecation, error on all others // see `scalac -Wconf:help` for details - "-Wconf:cat=deprecation:wv,any:e", + // since 2.13.15, "-Wconf:cat=deprecation:wv,any:e" no longer takes effect and needs to + // be changed to "-Wconf:any:e", "-Wconf:cat=deprecation:wv", + // please refer to the details: https://github.com/scala/scala/pull/10708 + "-Wconf:any:e", + "-Wconf:cat=deprecation:wv", // 2.13-specific warning hits to be muted (as narrowly as possible) and addressed separately "-Wunused:imports", "-Wconf:msg=^(?=.*?method|value|type|object|trait|inheritance)(?=.*?deprecated)(?=.*?since 2.13).+$:e", From 1f2e7b87db76ef60eded8a6db09f6690238471ce Mon Sep 17 00:00:00 2001 From: Enrico Minack Date: Wed, 25 Sep 2024 07:53:12 -0700 Subject: [PATCH 138/189] [SPARK-49731][K8S] Support K8s volume `mount.subPathExpr` and `hostPath` volume `type` ### What changes were proposed in this pull request? Add the following config options: - `spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].mount.subPathExpr` - `spark.kubernetes.executor.volumes.hostPath.[VolumeName].options.type` ### Why are the changes needed? K8s Spec - https://kubernetes.io/docs/concepts/storage/volumes/#hostpath-volume-types - https://kubernetes.io/docs/concepts/storage/volumes/#using-subpath-expanded-environment These are natural extensions of the existing options - `spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].mount.subPath` - `spark.kubernetes.executor.volumes.hostPath.[VolumeName].options.path` ### Does this PR introduce _any_ user-facing change? Above config options. ### How was this patch tested? Unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #48181 from EnricoMi/k8s-volume-options. Authored-by: Enrico Minack Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/deploy/k8s/Config.scala | 2 + .../deploy/k8s/KubernetesVolumeSpec.scala | 3 +- .../deploy/k8s/KubernetesVolumeUtils.scala | 18 ++++- .../features/MountVolumesFeatureStep.scala | 6 +- .../spark/deploy/k8s/KubernetesTestConf.scala | 11 ++- .../k8s/KubernetesVolumeUtilsSuite.scala | 42 ++++++++++- .../features/LocalDirsFeatureStepSuite.scala | 3 +- .../MountVolumesFeatureStepSuite.scala | 72 ++++++++++++++++++- 8 files changed, 144 insertions(+), 13 deletions(-) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala index 3a4d68c19014d..9c50f8ddb00cc 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala @@ -769,8 +769,10 @@ private[spark] object Config extends Logging { val KUBERNETES_VOLUMES_NFS_TYPE = "nfs" val KUBERNETES_VOLUMES_MOUNT_PATH_KEY = "mount.path" val KUBERNETES_VOLUMES_MOUNT_SUBPATH_KEY = "mount.subPath" + val KUBERNETES_VOLUMES_MOUNT_SUBPATHEXPR_KEY = "mount.subPathExpr" val KUBERNETES_VOLUMES_MOUNT_READONLY_KEY = "mount.readOnly" val KUBERNETES_VOLUMES_OPTIONS_PATH_KEY = "options.path" + val KUBERNETES_VOLUMES_OPTIONS_TYPE_KEY = "options.type" val KUBERNETES_VOLUMES_OPTIONS_CLAIM_NAME_KEY = "options.claimName" val KUBERNETES_VOLUMES_OPTIONS_CLAIM_STORAGE_CLASS_KEY = "options.storageClass" val KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY = "options.medium" diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala index 9dfd40a773eb1..b4fe414e3cde5 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala @@ -18,7 +18,7 @@ package org.apache.spark.deploy.k8s private[spark] sealed trait KubernetesVolumeSpecificConf -private[spark] case class KubernetesHostPathVolumeConf(hostPath: String) +private[spark] case class KubernetesHostPathVolumeConf(hostPath: String, volumeType: String) extends KubernetesVolumeSpecificConf private[spark] case class KubernetesPVCVolumeConf( @@ -42,5 +42,6 @@ private[spark] case class KubernetesVolumeSpec( volumeName: String, mountPath: String, mountSubPath: String, + mountSubPathExpr: String, mountReadOnly: Boolean, volumeConf: KubernetesVolumeSpecificConf) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala index 6463512c0114b..88bb998d88b7d 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala @@ -45,7 +45,9 @@ object KubernetesVolumeUtils { val pathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_PATH_KEY" val readOnlyKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_READONLY_KEY" val subPathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_SUBPATH_KEY" + val subPathExprKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_SUBPATHEXPR_KEY" val labelKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_LABEL_KEY" + verifyMutuallyExclusiveOptionKeys(properties, subPathKey, subPathExprKey) val volumeLabelsMap = properties .filter(_._1.startsWith(labelKey)) @@ -57,6 +59,7 @@ object KubernetesVolumeUtils { volumeName = volumeName, mountPath = properties(pathKey), mountSubPath = properties.getOrElse(subPathKey, ""), + mountSubPathExpr = properties.getOrElse(subPathExprKey, ""), mountReadOnly = properties.get(readOnlyKey).exists(_.toBoolean), volumeConf = parseVolumeSpecificConf(properties, volumeType, volumeName, Option(volumeLabelsMap))) @@ -87,8 +90,11 @@ object KubernetesVolumeUtils { volumeType match { case KUBERNETES_VOLUMES_HOSTPATH_TYPE => val pathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_PATH_KEY" + val typeKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_TYPE_KEY" verifyOptionKey(options, pathKey, KUBERNETES_VOLUMES_HOSTPATH_TYPE) - KubernetesHostPathVolumeConf(options(pathKey)) + // "" means that no checks will be performed before mounting the hostPath volume + // backward compatibility default + KubernetesHostPathVolumeConf(options(pathKey), options.getOrElse(typeKey, "")) case KUBERNETES_VOLUMES_PVC_TYPE => val claimNameKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_CLAIM_NAME_KEY" @@ -129,6 +135,16 @@ object KubernetesVolumeUtils { } } + private def verifyMutuallyExclusiveOptionKeys( + options: Map[String, String], + keys: String*): Unit = { + val givenKeys = keys.filter(options.contains) + if (givenKeys.length > 1) { + throw new IllegalArgumentException("These config options are mutually exclusive: " + + s"${givenKeys.mkString(", ")}") + } + } + private def verifySize(size: Option[String]): Unit = { size.foreach { v => if (v.forall(_.isDigit) && parseLong(v) < 1024) { diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala index 5cc61c746b0e0..eea4604010b21 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala @@ -65,14 +65,14 @@ private[spark] class MountVolumesFeatureStep(conf: KubernetesConf) .withMountPath(spec.mountPath) .withReadOnly(spec.mountReadOnly) .withSubPath(spec.mountSubPath) + .withSubPathExpr(spec.mountSubPathExpr) .withName(spec.volumeName) .build() val volumeBuilder = spec.volumeConf match { - case KubernetesHostPathVolumeConf(hostPath) => - /* "" means that no checks will be performed before mounting the hostPath volume */ + case KubernetesHostPathVolumeConf(hostPath, volumeType) => new VolumeBuilder() - .withHostPath(new HostPathVolumeSource(hostPath, "")) + .withHostPath(new HostPathVolumeSource(hostPath, volumeType)) case KubernetesPVCVolumeConf(claimNameTemplate, storageClass, size, labels) => val claimName = conf match { diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala index 7e0a65bcdda90..e0ddcd3d416f0 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala @@ -113,9 +113,10 @@ object KubernetesTestConf { volumes.foreach { case spec => val (vtype, configs) = spec.volumeConf match { - case KubernetesHostPathVolumeConf(path) => - (KUBERNETES_VOLUMES_HOSTPATH_TYPE, - Map(KUBERNETES_VOLUMES_OPTIONS_PATH_KEY -> path)) + case KubernetesHostPathVolumeConf(hostPath, volumeType) => + (KUBERNETES_VOLUMES_HOSTPATH_TYPE, Map( + KUBERNETES_VOLUMES_OPTIONS_PATH_KEY -> hostPath, + KUBERNETES_VOLUMES_OPTIONS_TYPE_KEY -> volumeType)) case KubernetesPVCVolumeConf(claimName, storageClass, sizeLimit, labels) => val sconf = storageClass @@ -145,6 +146,10 @@ object KubernetesTestConf { conf.set(key(vtype, spec.volumeName, KUBERNETES_VOLUMES_MOUNT_SUBPATH_KEY), spec.mountSubPath) } + if (spec.mountSubPathExpr.nonEmpty) { + conf.set(key(vtype, spec.volumeName, KUBERNETES_VOLUMES_MOUNT_SUBPATHEXPR_KEY), + spec.mountSubPathExpr) + } conf.set(key(vtype, spec.volumeName, KUBERNETES_VOLUMES_MOUNT_READONLY_KEY), spec.mountReadOnly.toString) configs.foreach { case (k, v) => diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala index 5c103739d3082..1e62db725fb6e 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala @@ -30,7 +30,20 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.mountPath === "/path") assert(volumeSpec.mountReadOnly) assert(volumeSpec.volumeConf.asInstanceOf[KubernetesHostPathVolumeConf] === - KubernetesHostPathVolumeConf("/hostPath")) + KubernetesHostPathVolumeConf("/hostPath", "")) + } + + test("Parses hostPath volume type correctly") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.hostPath.volumeName.mount.path", "/path") + sparkConf.set("test.hostPath.volumeName.options.path", "/hostPath") + sparkConf.set("test.hostPath.volumeName.options.type", "Type") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.volumeConf.asInstanceOf[KubernetesHostPathVolumeConf] === + KubernetesHostPathVolumeConf("/hostPath", "Type")) } test("Parses subPath correctly") { @@ -43,6 +56,33 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.volumeName === "volumeName") assert(volumeSpec.mountPath === "/path") assert(volumeSpec.mountSubPath === "subPath") + assert(volumeSpec.mountSubPathExpr === "") + } + + test("Parses subPathExpr correctly") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.emptyDir.volumeName.mount.path", "/path") + sparkConf.set("test.emptyDir.volumeName.mount.readOnly", "true") + sparkConf.set("test.emptyDir.volumeName.mount.subPathExpr", "subPathExpr") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.mountSubPath === "") + assert(volumeSpec.mountSubPathExpr === "subPathExpr") + } + + test("Rejects mutually exclusive subPath and subPathExpr") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.emptyDir.volumeName.mount.path", "/path") + sparkConf.set("test.emptyDir.volumeName.mount.subPath", "subPath") + sparkConf.set("test.emptyDir.volumeName.mount.subPathExpr", "subPathExpr") + + val msg = intercept[IllegalArgumentException] { + KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + }.getMessage + assert(msg === "These config options are mutually exclusive: " + + "emptyDir.volumeName.mount.subPath, emptyDir.volumeName.mount.subPathExpr") } test("Parses persistentVolumeClaim volumes correctly") { diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala index eaadad163f064..3a9561051a894 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala @@ -137,8 +137,9 @@ class LocalDirsFeatureStepSuite extends SparkFunSuite { "spark-local-dir-test", "/tmp", "", + "", false, - KubernetesHostPathVolumeConf("/hostPath/tmp") + KubernetesHostPathVolumeConf("/hostPath/tmp", "") ) val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) val mountVolumeStep = new MountVolumesFeatureStep(kubernetesConf) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala index 6a68898c5f61c..c94a7a6ec26a7 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala @@ -27,8 +27,9 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", false, - KubernetesHostPathVolumeConf("/hostPath/tmp") + KubernetesHostPathVolumeConf("/hostPath/tmp", "type") ) val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) val step = new MountVolumesFeatureStep(kubernetesConf) @@ -36,6 +37,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { assert(configuredPod.pod.getSpec.getVolumes.size() === 1) assert(configuredPod.pod.getSpec.getVolumes.get(0).getHostPath.getPath === "/hostPath/tmp") + assert(configuredPod.pod.getSpec.getVolumes.get(0).getHostPath.getType === "type") assert(configuredPod.container.getVolumeMounts.size() === 1) assert(configuredPod.container.getVolumeMounts.get(0).getMountPath === "/tmp") assert(configuredPod.container.getVolumeMounts.get(0).getName === "testVolume") @@ -47,6 +49,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf("pvcClaim") ) @@ -69,6 +72,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf("pvc-spark-SPARK_EXECUTOR_ID") ) @@ -94,6 +98,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf("pvc-spark-SPARK_EXECUTOR_ID", Some("fast"), Some("512M")) ) @@ -119,6 +124,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf("OnDemand") ) @@ -136,6 +142,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, storageClass = Some("gp3"), @@ -156,6 +163,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, storageClass = Some("gp3"), @@ -177,6 +185,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "checkpointVolume1", "/checkpoints1", "", + "", true, KubernetesPVCVolumeConf(claimName = "pvcClaim1", storageClass = Some("gp3"), @@ -188,6 +197,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "checkpointVolume2", "/checkpoints2", "", + "", true, KubernetesPVCVolumeConf(claimName = "pvcClaim2", storageClass = Some("gp3"), @@ -209,6 +219,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesPVCVolumeConf(MountVolumesFeatureStep.PVC_ON_DEMAND) ) @@ -226,6 +237,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", false, KubernetesEmptyDirVolumeConf(Some("Memory"), Some("6G")) ) @@ -249,6 +261,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", false, KubernetesEmptyDirVolumeConf(None, None) ) @@ -271,6 +284,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", false, KubernetesNFSVolumeConf("/share/name", "nfs.example.com") ) @@ -293,6 +307,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "", + "", true, KubernetesNFSVolumeConf("/share/name", "nfs.example.com") ) @@ -315,13 +330,15 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "hpVolume", "/tmp", "", + "", false, - KubernetesHostPathVolumeConf("/hostPath/tmp") + KubernetesHostPathVolumeConf("/hostPath/tmp", "") ) val pvcVolumeConf = KubernetesVolumeSpec( "checkpointVolume", "/checkpoints", "", + "", true, KubernetesPVCVolumeConf("pvcClaim") ) @@ -339,13 +356,15 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "hpVolume", "/data", "", + "", false, - KubernetesHostPathVolumeConf("/hostPath/tmp") + KubernetesHostPathVolumeConf("/hostPath/tmp", "") ) val pvcVolumeConf = KubernetesVolumeSpec( "checkpointVolume", "/data", "", + "", true, KubernetesPVCVolumeConf("pvcClaim") ) @@ -364,6 +383,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testVolume", "/tmp", "foo", + "", false, KubernetesEmptyDirVolumeConf(None, None) ) @@ -378,11 +398,32 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { assert(emptyDirMount.getSubPath === "foo") } + test("Mounts subpathexpr on emptyDir") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + "foo", + false, + KubernetesEmptyDirVolumeConf(None, None) + ) + val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + + assert(configuredPod.pod.getSpec.getVolumes.size() === 1) + val emptyDirMount = configuredPod.container.getVolumeMounts.get(0) + assert(emptyDirMount.getMountPath === "/tmp") + assert(emptyDirMount.getName === "testVolume") + assert(emptyDirMount.getSubPathExpr === "foo") + } + test("Mounts subpath on persistentVolumeClaims") { val volumeConf = KubernetesVolumeSpec( "testVolume", "/tmp", "bar", + "", true, KubernetesPVCVolumeConf("pvcClaim") ) @@ -400,12 +441,36 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { assert(pvcMount.getSubPath === "bar") } + test("Mounts subpathexpr on persistentVolumeClaims") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + "bar", + true, + KubernetesPVCVolumeConf("pvcClaim") + ) + val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + + assert(configuredPod.pod.getSpec.getVolumes.size() === 1) + val pvcClaim = configuredPod.pod.getSpec.getVolumes.get(0).getPersistentVolumeClaim + assert(pvcClaim.getClaimName === "pvcClaim") + assert(configuredPod.container.getVolumeMounts.size() === 1) + val pvcMount = configuredPod.container.getVolumeMounts.get(0) + assert(pvcMount.getMountPath === "/tmp") + assert(pvcMount.getName === "testVolume") + assert(pvcMount.getSubPathExpr === "bar") + } + test("Mounts multiple subpaths") { val volumeConf = KubernetesEmptyDirVolumeConf(None, None) val emptyDirSpec = KubernetesVolumeSpec( "testEmptyDir", "/tmp/foo", "foo", + "", true, KubernetesEmptyDirVolumeConf(None, None) ) @@ -413,6 +478,7 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { "testPVC", "/tmp/bar", "bar", + "", true, KubernetesEmptyDirVolumeConf(None, None) ) From 09209f0ff503b29f9da92ba7db8aa820c03b3c0f Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 25 Sep 2024 07:57:08 -0700 Subject: [PATCH 139/189] [SPARK-49775][SQL][FOLLOW-UP] Use SortedSet instead of Array with sorting ### What changes were proposed in this pull request? This PR is a followup of https://github.com/apache/spark/pull/48235 that addresses https://github.com/apache/spark/pull/48235#discussion_r1775020195 comment. ### Why are the changes needed? For better performance (in theory) ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests should verify them ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48245 from HyukjinKwon/SPARK-49775-followup. Authored-by: Hyukjin Kwon Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/sql/catalyst/util/CharsetProvider.scala | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala index d85673f2ce811..f805d2ed87b52 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharsetProvider.scala @@ -18,13 +18,15 @@ import java.nio.charset.{Charset, CharsetDecoder, CharsetEncoder, CodingErrorAction, IllegalCharsetNameException, UnsupportedCharsetException} import java.util.Locale + import scala.collection.SortedSet + import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf private[sql] object CharsetProvider { final lazy val VALID_CHARSETS = - Array("us-ascii", "iso-8859-1", "utf-8", "utf-16be", "utf-16le", "utf-16", "utf-32").sorted + SortedSet("us-ascii", "iso-8859-1", "utf-8", "utf-16be", "utf-16le", "utf-16", "utf-32") def forName( charset: String, From 80d6651cf6a1835d0de3e12e08253d2a9816d499 Mon Sep 17 00:00:00 2001 From: Julek Sompolski Date: Wed, 25 Sep 2024 23:34:23 +0800 Subject: [PATCH 140/189] [SPARK-48195][FOLLOWUP] Accumulator reset() no longer needed in CollectMetricsExec.doExecute() ### What changes were proposed in this pull request? Small followup to https://github.com/apache/spark/pull/48037. `collector.reset()` is no longer needed in `CollectMetricsExec.doExecute()` because it is reset in `resetMetrics()`. This doesn't really matter in practice, but removing to clean up. ### Why are the changes needed? Tiny cleanup. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? This change doesn't matter in practice. Just cleanup. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48243 from juliuszsompolski/SPARK-48195-followup. Authored-by: Julek Sompolski Signed-off-by: Wenchen Fan --- .../org/apache/spark/sql/execution/CollectMetricsExec.scala | 1 - 1 file changed, 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala index 2115e21f81d71..0a487bac77696 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CollectMetricsExec.scala @@ -67,7 +67,6 @@ case class CollectMetricsExec( override protected def doExecute(): RDD[InternalRow] = { val collector = accumulator - collector.reset() child.execute().mapPartitions { rows => // Only publish the value of the accumulator when the task has completed. This is done by // updating a task local accumulator ('updater') which will be merged with the actual From c0984e70469d99595b8e6eda0d943308f590aaec Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 26 Sep 2024 13:17:59 +0900 Subject: [PATCH 141/189] [SPARK-49609][PYTHON][TESTS][FOLLOW-UP] Avoid import connect modules when connect dependencies not installed ### What changes were proposed in this pull request? This PR is a followup of https://github.com/apache/spark/pull/48085 that skips the connect import which requires Connect dependencies. ### Why are the changes needed? To recover the PyPy3 build https://github.com/apache/spark/actions/runs/11035779484/job/30652736098 which does not have PyArrow installed. ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? Manually. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48259 from HyukjinKwon/SPARK-49609-followup2. Authored-by: Hyukjin Kwon Signed-off-by: Hyukjin Kwon --- python/pyspark/sql/tests/test_connect_compatibility.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/tests/test_connect_compatibility.py b/python/pyspark/sql/tests/test_connect_compatibility.py index 8f3e86f5186a8..dfa0fa63b2dd5 100644 --- a/python/pyspark/sql/tests/test_connect_compatibility.py +++ b/python/pyspark/sql/tests/test_connect_compatibility.py @@ -21,11 +21,13 @@ from pyspark.testing.connectutils import should_test_connect, connect_requirement_message from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.classic.dataframe import DataFrame as ClassicDataFrame -from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame from pyspark.sql.classic.column import Column as ClassicColumn -from pyspark.sql.connect.column import Column as ConnectColumn from pyspark.sql.session import SparkSession as ClassicSparkSession -from pyspark.sql.connect.session import SparkSession as ConnectSparkSession + +if should_test_connect: + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.session import SparkSession as ConnectSparkSession class ConnectCompatibilityTestsMixin: From 5629779287724a891c81b16f982f9529bd379c39 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Wed, 25 Sep 2024 22:34:35 -0700 Subject: [PATCH 142/189] [SPARK-49786][K8S] Lower `KubernetesClusterSchedulerBackend.onDisconnected` log level to debug ### What changes were proposed in this pull request? This PR aims to lower `KubernetesClusterSchedulerBackend.onDisconnected` log level to debug. ### Why are the changes needed? This INFO-level message was added here. We already propagate the disconnection reason to UI, and `No executor found` has been used when an unknown peer is connect or disconnect. - https://github.com/apache/spark/pull/37821 The driver can be accessed by non-executors by design. And, all other resource managers do not complain at INFO level. ``` INFO KubernetesClusterSchedulerBackend$KubernetesDriverEndpoint: No executor found for x.x.x.0:x ``` ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Manual review because this is a log level change. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48249 from dongjoon-hyun/SPARK-49786. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../cluster/k8s/KubernetesClusterSchedulerBackend.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala index 4e4634504a0f3..09faa2a7fb1b3 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala @@ -32,7 +32,7 @@ import org.apache.spark.deploy.k8s.Config._ import org.apache.spark.deploy.k8s.Constants._ import org.apache.spark.deploy.k8s.submit.KubernetesClientUtils import org.apache.spark.deploy.security.HadoopDelegationTokenManager -import org.apache.spark.internal.LogKeys.{COUNT, HOST_PORT, TOTAL} +import org.apache.spark.internal.LogKeys.{COUNT, TOTAL} import org.apache.spark.internal.MDC import org.apache.spark.internal.config.SCHEDULER_MIN_REGISTERED_RESOURCES_RATIO import org.apache.spark.resource.ResourceProfile @@ -356,7 +356,7 @@ private[spark] class KubernetesClusterSchedulerBackend( execIDRequester -= rpcAddress // Expected, executors re-establish a connection with an ID case _ => - logInfo(log"No executor found for ${MDC(HOST_PORT, rpcAddress)}") + logDebug(s"No executor found for ${rpcAddress}") } } } From 913a0f7813c5b2d2bf105160bf8e55e08b34513b Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 26 Sep 2024 15:15:37 +0800 Subject: [PATCH 143/189] [SPARK-49784][PYTHON][TESTS] Add more test for `spark.sql` ### What changes were proposed in this pull request? add more test for `spark.sql` ### Why are the changes needed? for test coverage ### Does this PR introduce _any_ user-facing change? no, test only ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48246 from zhengruifeng/py_sql_test. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- dev/sparktestsupport/modules.py | 2 + .../sql/tests/connect/test_parity_sql.py | 37 ++++ python/pyspark/sql/tests/test_sql.py | 185 ++++++++++++++++++ 3 files changed, 224 insertions(+) create mode 100644 python/pyspark/sql/tests/connect/test_parity_sql.py create mode 100644 python/pyspark/sql/tests/test_sql.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index eda6b063350e5..d2c000b702a64 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -520,6 +520,7 @@ def __hash__(self): "pyspark.sql.tests.test_errors", "pyspark.sql.tests.test_functions", "pyspark.sql.tests.test_group", + "pyspark.sql.tests.test_sql", "pyspark.sql.tests.pandas.test_pandas_cogrouped_map", "pyspark.sql.tests.pandas.test_pandas_grouped_map", "pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state", @@ -1032,6 +1033,7 @@ def __hash__(self): "pyspark.sql.tests.connect.test_parity_serde", "pyspark.sql.tests.connect.test_parity_functions", "pyspark.sql.tests.connect.test_parity_group", + "pyspark.sql.tests.connect.test_parity_sql", "pyspark.sql.tests.connect.test_parity_dataframe", "pyspark.sql.tests.connect.test_parity_collection", "pyspark.sql.tests.connect.test_parity_creation", diff --git a/python/pyspark/sql/tests/connect/test_parity_sql.py b/python/pyspark/sql/tests/connect/test_parity_sql.py new file mode 100644 index 0000000000000..4c6b11c60cbe9 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_sql.py @@ -0,0 +1,37 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest + +from pyspark.sql.tests.test_sql import SQLTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class SQLParityTests(SQLTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.connect.test_parity_sql import * # noqa: F401 + + try: + import xmlrunner # type: ignore[import] + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) diff --git a/python/pyspark/sql/tests/test_sql.py b/python/pyspark/sql/tests/test_sql.py new file mode 100644 index 0000000000000..bf50bbc11ac33 --- /dev/null +++ b/python/pyspark/sql/tests/test_sql.py @@ -0,0 +1,185 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest + +from pyspark.sql import Row +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class SQLTestsMixin: + def test_simple(self): + res = self.spark.sql("SELECT 1 + 1").collect() + self.assertEqual(len(res), 1) + self.assertEqual(res[0][0], 2) + + def test_args_dict(self): + with self.tempView("test"): + self.spark.range(10).createOrReplaceTempView("test") + df = self.spark.sql( + "SELECT * FROM IDENTIFIER(:table_name)", + args={"table_name": "test"}, + ) + + self.assertEqual(df.count(), 10) + self.assertEqual(df.limit(5).count(), 5) + self.assertEqual(df.offset(5).count(), 5) + + self.assertEqual(df.take(1), [Row(id=0)]) + self.assertEqual(df.tail(1), [Row(id=9)]) + + def test_args_list(self): + with self.tempView("test"): + self.spark.range(10).createOrReplaceTempView("test") + df = self.spark.sql( + "SELECT * FROM test WHERE ? < id AND id < ?", + args=[1, 6], + ) + + self.assertEqual(df.count(), 4) + self.assertEqual(df.limit(3).count(), 3) + self.assertEqual(df.offset(3).count(), 1) + + self.assertEqual(df.take(1), [Row(id=2)]) + self.assertEqual(df.tail(1), [Row(id=5)]) + + def test_kwargs_literal(self): + with self.tempView("test"): + self.spark.range(10).createOrReplaceTempView("test") + + df = self.spark.sql( + "SELECT * FROM IDENTIFIER(:table_name) WHERE {m1} < id AND id < {m2} OR id = {m3}", + args={"table_name": "test"}, + m1=3, + m2=7, + m3=9, + ) + + self.assertEqual(df.count(), 4) + self.assertEqual(df.collect(), [Row(id=4), Row(id=5), Row(id=6), Row(id=9)]) + self.assertEqual(df.take(1), [Row(id=4)]) + self.assertEqual(df.tail(1), [Row(id=9)]) + + def test_kwargs_literal_multiple_ref(self): + with self.tempView("test"): + self.spark.range(10).createOrReplaceTempView("test") + + df = self.spark.sql( + "SELECT * FROM IDENTIFIER(:table_name) WHERE {m} = id OR id > {m} OR {m} < 0", + args={"table_name": "test"}, + m=6, + ) + + self.assertEqual(df.count(), 4) + self.assertEqual(df.collect(), [Row(id=6), Row(id=7), Row(id=8), Row(id=9)]) + self.assertEqual(df.take(1), [Row(id=6)]) + self.assertEqual(df.tail(1), [Row(id=9)]) + + def test_kwargs_dataframe(self): + df0 = self.spark.range(10) + df1 = self.spark.sql( + "SELECT * FROM {df} WHERE id > 4", + df=df0, + ) + + self.assertEqual(df0.schema, df1.schema) + self.assertEqual(df1.count(), 5) + self.assertEqual(df1.take(1), [Row(id=5)]) + self.assertEqual(df1.tail(1), [Row(id=9)]) + + def test_kwargs_dataframe_with_column(self): + df0 = self.spark.range(10) + df1 = self.spark.sql( + "SELECT * FROM {df} WHERE {df.id} > :m1 AND {df[id]} < :m2", + {"m1": 4, "m2": 9}, + df=df0, + ) + + self.assertEqual(df0.schema, df1.schema) + self.assertEqual(df1.count(), 4) + self.assertEqual(df1.take(1), [Row(id=5)]) + self.assertEqual(df1.tail(1), [Row(id=8)]) + + def test_nested_view(self): + with self.tempView("v1", "v2", "v3", "v4"): + self.spark.range(10).createOrReplaceTempView("v1") + self.spark.sql( + "SELECT * FROM IDENTIFIER(:view) WHERE id > :m", + args={"view": "v1", "m": 1}, + ).createOrReplaceTempView("v2") + self.spark.sql( + "SELECT * FROM IDENTIFIER(:view) WHERE id > :m", + args={"view": "v2", "m": 2}, + ).createOrReplaceTempView("v3") + self.spark.sql( + "SELECT * FROM IDENTIFIER(:view) WHERE id > :m", + args={"view": "v3", "m": 3}, + ).createOrReplaceTempView("v4") + + df = self.spark.sql("select * from v4") + self.assertEqual(df.count(), 6) + self.assertEqual(df.take(1), [Row(id=4)]) + self.assertEqual(df.tail(1), [Row(id=9)]) + + def test_nested_dataframe(self): + df0 = self.spark.range(10) + df1 = self.spark.sql( + "SELECT * FROM {df} WHERE id > ?", + args=[1], + df=df0, + ) + df2 = self.spark.sql( + "SELECT * FROM {df} WHERE id > ?", + args=[2], + df=df1, + ) + df3 = self.spark.sql( + "SELECT * FROM {df} WHERE id > ?", + args=[3], + df=df2, + ) + + self.assertEqual(df0.schema, df1.schema) + self.assertEqual(df1.count(), 8) + self.assertEqual(df1.take(1), [Row(id=2)]) + self.assertEqual(df1.tail(1), [Row(id=9)]) + + self.assertEqual(df0.schema, df2.schema) + self.assertEqual(df2.count(), 7) + self.assertEqual(df2.take(1), [Row(id=3)]) + self.assertEqual(df2.tail(1), [Row(id=9)]) + + self.assertEqual(df0.schema, df3.schema) + self.assertEqual(df3.count(), 6) + self.assertEqual(df3.take(1), [Row(id=4)]) + self.assertEqual(df3.tail(1), [Row(id=9)]) + + +class SQLTests(SQLTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.sql.tests.test_sql import * # noqa: F401 + + try: + import xmlrunner # type: ignore + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + unittest.main(testRunner=testRunner, verbosity=2) From fe1cf3200223c33ed4670bfa5924d5a4053c8ef9 Mon Sep 17 00:00:00 2001 From: Anish Shrigondekar Date: Thu, 26 Sep 2024 17:38:58 +0900 Subject: [PATCH 144/189] [SPARK-49656][SS] Add support for state variables with value state collection types and read change feed options ### What changes were proposed in this pull request? Add support for state variables with value state collection types and read change feed options ### Why are the changes needed? Without this, we cannot support reading per key changes for state variables used with stateful processors. ### Does this PR introduce _any_ user-facing change? Yes Users can now query value state variables with the following query: ``` val changeFeedDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, ) .option(StateSourceOptions.STATE_VAR_NAME, ) .option(StateSourceOptions.READ_CHANGE_FEED, true) .option(StateSourceOptions.CHANGE_START_BATCH_ID, 0) .load() ``` ### How was this patch tested? Added unit tests ``` [info] Run completed in 17 seconds, 318 milliseconds. [info] Total number of tests run: 2 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 2, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48148 from anishshri-db/task/SPARK-49656. Authored-by: Anish Shrigondekar Signed-off-by: Jungtaek Lim --- .../v2/state/StateDataSource.scala | 10 +- .../v2/state/StatePartitionReader.scala | 10 +- .../state/HDFSBackedStateStoreProvider.scala | 10 +- .../state/RocksDBStateStoreProvider.scala | 79 ++++++++++--- .../streaming/state/StateStore.scala | 6 +- .../streaming/state/StateStoreChangelog.scala | 11 +- ...ateDataSourceTransformWithStateSuite.scala | 107 +++++++++++++++--- 7 files changed, 190 insertions(+), 43 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala index 39bc4dd9fb9c8..edddfbd6ccaef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.{J import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions.JoinSideValues.JoinSideValues import org.apache.spark.sql.execution.datasources.v2.state.metadata.{StateMetadataPartitionReader, StateMetadataTableEntry} import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil -import org.apache.spark.sql.execution.streaming.{CommitLog, OffsetSeqLog, OffsetSeqMetadata, TimerStateUtils, TransformWithStateOperatorProperties, TransformWithStateVariableInfo} +import org.apache.spark.sql.execution.streaming.{CommitLog, OffsetSeqLog, OffsetSeqMetadata, StateVariableType, TimerStateUtils, TransformWithStateOperatorProperties, TransformWithStateVariableInfo} import org.apache.spark.sql.execution.streaming.StreamingCheckpointConstants.{DIR_NAME_COMMITS, DIR_NAME_OFFSETS, DIR_NAME_STATE} import org.apache.spark.sql.execution.streaming.StreamingSymmetricHashJoinHelper.{LeftSide, RightSide} import org.apache.spark.sql.execution.streaming.state.{KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, StateSchemaCompatibilityChecker, StateStore, StateStoreColFamilySchema, StateStoreConf, StateStoreId, StateStoreProviderId} @@ -170,13 +170,15 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging } val stateVars = twsOperatorProperties.stateVariables - if (stateVars.filter(stateVar => stateVar.stateName == stateVarName).size != 1) { + val stateVarInfo = stateVars.filter(stateVar => stateVar.stateName == stateVarName) + if (stateVarInfo.size != 1) { throw StateDataSourceErrors.invalidOptionValue(STATE_VAR_NAME, s"State variable $stateVarName is not defined for the transformWithState operator.") } - // TODO: Support change feed and transformWithState together - if (sourceOptions.readChangeFeed) { + // TODO: add support for list and map type + if (sourceOptions.readChangeFeed && + stateVarInfo.head.stateVariableType != StateVariableType.ValueState) { throw StateDataSourceErrors.conflictOptions(Seq(StateSourceOptions.READ_CHANGE_FEED, StateSourceOptions.STATE_VAR_NAME)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala index d77d97f0057fb..b925aee5b627a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala @@ -223,10 +223,18 @@ class StateStoreChangeDataPartitionReader( throw StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay( provider.getClass.toString) } + + val colFamilyNameOpt = if (stateVariableInfoOpt.isDefined) { + Some(stateVariableInfoOpt.get.stateName) + } else { + None + } + provider.asInstanceOf[SupportsFineGrainedReplay] .getStateStoreChangeDataReader( partition.sourceOptions.readChangeFeedOptions.get.changeStartBatchId + 1, - partition.sourceOptions.readChangeFeedOptions.get.changeEndBatchId + 1) + partition.sourceOptions.readChangeFeedOptions.get.changeEndBatchId + 1, + colFamilyNameOpt) } override lazy val iter: Iterator[InternalRow] = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala index d9f4443b79618..884b8aa3853cb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala @@ -991,8 +991,16 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with result } - override def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long): + override def getStateStoreChangeDataReader( + startVersion: Long, + endVersion: Long, + colFamilyNameOpt: Option[String] = None): StateStoreChangeDataReader = { + // Multiple column families are not supported with HDFSBackedStateStoreProvider + if (colFamilyNameOpt.isDefined) { + throw StateStoreErrors.multipleColumnFamiliesNotSupported(providerName) + } + new HDFSBackedStateStoreChangeDataReader(fm, baseDir, startVersion, endVersion, CompressionCodec.createCodec(sparkConf, storeConf.compressionCodec), keySchema, valueSchema) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala index 85f80ce9eb1ae..6ab634668bc2a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala @@ -498,7 +498,10 @@ private[sql] class RocksDBStateStoreProvider } } - override def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long): + override def getStateStoreChangeDataReader( + startVersion: Long, + endVersion: Long, + colFamilyNameOpt: Option[String] = None): StateStoreChangeDataReader = { val statePath = stateStoreId.storeCheckpointLocation() val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf) @@ -508,7 +511,8 @@ private[sql] class RocksDBStateStoreProvider startVersion, endVersion, CompressionCodec.createCodec(sparkConf, storeConf.compressionCodec), - keyValueEncoderMap) + keyValueEncoderMap, + colFamilyNameOpt) } /** @@ -676,27 +680,70 @@ class RocksDBStateStoreChangeDataReader( endVersion: Long, compressionCodec: CompressionCodec, keyValueEncoderMap: - ConcurrentHashMap[String, (RocksDBKeyStateEncoder, RocksDBValueStateEncoder)]) + ConcurrentHashMap[String, (RocksDBKeyStateEncoder, RocksDBValueStateEncoder)], + colFamilyNameOpt: Option[String] = None) extends StateStoreChangeDataReader( - fm, stateLocation, startVersion, endVersion, compressionCodec) { + fm, stateLocation, startVersion, endVersion, compressionCodec, colFamilyNameOpt) { override protected var changelogSuffix: String = "changelog" + private def getColFamilyIdBytes: Option[Array[Byte]] = { + if (colFamilyNameOpt.isDefined) { + val colFamilyName = colFamilyNameOpt.get + if (!keyValueEncoderMap.containsKey(colFamilyName)) { + throw new IllegalStateException( + s"Column family $colFamilyName not found in the key value encoder map") + } + Some(keyValueEncoderMap.get(colFamilyName)._1.getColumnFamilyIdBytes()) + } else { + None + } + } + + private val colFamilyIdBytesOpt: Option[Array[Byte]] = getColFamilyIdBytes + override def getNext(): (RecordType.Value, UnsafeRow, UnsafeRow, Long) = { - val reader = currentChangelogReader() - if (reader == null) { - return null + var currRecord: (RecordType.Value, Array[Byte], Array[Byte]) = null + val currEncoder: (RocksDBKeyStateEncoder, RocksDBValueStateEncoder) = + keyValueEncoderMap.get(colFamilyNameOpt + .getOrElse(StateStore.DEFAULT_COL_FAMILY_NAME)) + + if (colFamilyIdBytesOpt.isDefined) { + // If we are reading records for a particular column family, the corresponding vcf id + // will be encoded in the key byte array. We need to extract that and compare for the + // expected column family id. If it matches, we return the record. If not, we move to + // the next record. Note that this has be handled across multiple changelog files and we + // rely on the currentChangelogReader to move to the next changelog file when needed. + while (currRecord == null) { + val reader = currentChangelogReader() + if (reader == null) { + return null + } + + val nextRecord = reader.next() + val colFamilyIdBytes: Array[Byte] = colFamilyIdBytesOpt.get + val endIndex = colFamilyIdBytes.size + // Function checks for byte arrays being equal + // from index 0 to endIndex - 1 (both inclusive) + if (java.util.Arrays.equals(nextRecord._2, 0, endIndex, + colFamilyIdBytes, 0, endIndex)) { + currRecord = nextRecord + } + } + } else { + val reader = currentChangelogReader() + if (reader == null) { + return null + } + currRecord = reader.next() } - val (recordType, keyArray, valueArray) = reader.next() - // Todo: does not support multiple virtual column families - val (rocksDBKeyStateEncoder, rocksDBValueStateEncoder) = - keyValueEncoderMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) - val keyRow = rocksDBKeyStateEncoder.decodeKey(keyArray) - if (valueArray == null) { - (recordType, keyRow, null, currentChangelogVersion - 1) + + val keyRow = currEncoder._1.decodeKey(currRecord._2) + if (currRecord._3 == null) { + (currRecord._1, keyRow, null, currentChangelogVersion - 1) } else { - val valueRow = rocksDBValueStateEncoder.decodeValue(valueArray) - (recordType, keyRow, valueRow, currentChangelogVersion - 1) + val valueRow = currEncoder._2.decodeValue(currRecord._3) + (currRecord._1, keyRow, valueRow, currentChangelogVersion - 1) } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala index d55a973a14e16..6e616cc71a80c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala @@ -519,10 +519,14 @@ trait SupportsFineGrainedReplay { * * @param startVersion starting changelog version * @param endVersion ending changelog version + * @param colFamilyNameOpt optional column family name to read from * @return iterator that gives tuple(recordType: [[RecordType.Value]], nested key: [[UnsafeRow]], * nested value: [[UnsafeRow]], batchId: [[Long]]) */ - def getStateStoreChangeDataReader(startVersion: Long, endVersion: Long): + def getStateStoreChangeDataReader( + startVersion: Long, + endVersion: Long, + colFamilyNameOpt: Option[String] = None): NextIterator[(RecordType.Value, UnsafeRow, UnsafeRow, Long)] } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreChangelog.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreChangelog.scala index 651d72da16095..e89550da37e03 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreChangelog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreChangelog.scala @@ -397,13 +397,15 @@ class StateStoreChangelogReaderV2( * @param startVersion start version of the changelog file to read * @param endVersion end version of the changelog file to read * @param compressionCodec de-compression method using for reading changelog file + * @param colFamilyNameOpt optional column family name to read from */ abstract class StateStoreChangeDataReader( fm: CheckpointFileManager, stateLocation: Path, startVersion: Long, endVersion: Long, - compressionCodec: CompressionCodec) + compressionCodec: CompressionCodec, + colFamilyNameOpt: Option[String] = None) extends NextIterator[(RecordType.Value, UnsafeRow, UnsafeRow, Long)] with Logging { assert(startVersion >= 1) @@ -451,9 +453,12 @@ abstract class StateStoreChangeDataReader( finished = true return null } - // Todo: Does not support StateStoreChangelogReaderV2 - changelogReader = + + changelogReader = if (colFamilyNameOpt.isDefined) { + new StateStoreChangelogReaderV2(fm, fileIterator.next(), compressionCodec) + } else { new StateStoreChangelogReaderV1(fm, fileIterator.next(), compressionCodec) + } } changelogReader } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala index bd047d1132fbe..84c6eb54681a1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala @@ -186,18 +186,49 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest } assert(ex1.isInstanceOf[StateDataSourceInvalidOptionValue]) assert(ex1.getMessage.contains("Registered timers are not available")) + } + } + } - // TODO: this should be removed when readChangeFeed is supported for value state - val ex2 = intercept[Exception] { - spark.read + testWithChangelogCheckpointingEnabled("state data source cdf integration - " + + "value state with single variable") { + withTempDir { tempDir => + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString) { + val inputData = MemoryStream[String] + val result = inputData.toDS() + .groupByKey(x => x) + .transformWithState(new StatefulProcessorWithSingleValueVar(), + TimeMode.None(), + OutputMode.Update()) + + testStream(result, OutputMode.Update())( + StartStream(checkpointLocation = tempDir.getAbsolutePath), + AddData(inputData, "a"), + CheckNewAnswer(("a", "1")), + AddData(inputData, "b"), + CheckNewAnswer(("b", "1")), + StopStream + ) + + val changeFeedDf = spark.read .format("statestore") .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) .option(StateSourceOptions.STATE_VAR_NAME, "valueState") - .option(StateSourceOptions.READ_CHANGE_FEED, "true") + .option(StateSourceOptions.READ_CHANGE_FEED, true) .option(StateSourceOptions.CHANGE_START_BATCH_ID, 0) .load() - } - assert(ex2.isInstanceOf[StateDataSourceConflictOptions]) + + val opDf = changeFeedDf.selectExpr( + "change_type", + "key.value AS groupingKey", + "value.id AS valueId", "value.name AS valueName", + "partition_id") + + checkAnswer(opDf, + Seq(Row("update", "a", 1L, "dummyKey", 0), Row("update", "b", 1L, "dummyKey", 1))) } } } @@ -260,19 +291,61 @@ class StateDataSourceTransformWithStateSuite extends StateStoreMetricsTest } assert(ex.isInstanceOf[StateDataSourceInvalidOptionValue]) assert(ex.getMessage.contains("State variable non-exist is not defined")) + } + } + } - // TODO: this should be removed when readChangeFeed is supported for TTL based state - // variables - val ex1 = intercept[Exception] { - spark.read - .format("statestore") - .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) - .option(StateSourceOptions.STATE_VAR_NAME, "countState") - .option(StateSourceOptions.READ_CHANGE_FEED, "true") - .option(StateSourceOptions.CHANGE_START_BATCH_ID, 0) - .load() + testWithChangelogCheckpointingEnabled("state data source cdf integration - " + + "value state with single variable and TTL") { + withTempDir { tempDir => + withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> + TransformWithStateSuiteUtils.NUM_SHUFFLE_PARTITIONS.toString) { + val inputData = MemoryStream[String] + val result = inputData.toDS() + .groupByKey(x => x) + .transformWithState(new StatefulProcessorWithTTL(), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update())( + StartStream(checkpointLocation = tempDir.getAbsolutePath), + AddData(inputData, "a"), + AddData(inputData, "b"), + Execute { _ => + // wait for the batch to run since we are using processing time + Thread.sleep(5000) + }, + StopStream + ) + + val stateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, tempDir.getAbsolutePath) + .option(StateSourceOptions.STATE_VAR_NAME, "countState") + .option(StateSourceOptions.READ_CHANGE_FEED, true) + .option(StateSourceOptions.CHANGE_START_BATCH_ID, 0) + .load() + + val resultDf = stateReaderDf.selectExpr( + "key.value", "value.value", "value.ttlExpirationMs", "partition_id") + + var count = 0L + resultDf.collect().foreach { row => + count = count + 1 + assert(row.getLong(2) > 0) } - assert(ex1.isInstanceOf[StateDataSourceConflictOptions]) + + // verify that 2 state rows are present + assert(count === 2) + + val answerDf = stateReaderDf.selectExpr( + "change_type", + "key.value AS groupingKey", + "value.value.value AS valueId", "partition_id") + checkAnswer(answerDf, + Seq(Row("update", "a", 1L, 0), Row("update", "b", 1L, 1))) } } } From a116a5bf708dbd2e0efc0b1f63f3f655d3e830da Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Thu, 26 Sep 2024 08:37:04 -0400 Subject: [PATCH 145/189] [SPARK-49416][CONNECT][SQL] Add Shared DataStreamReader interface ### What changes were proposed in this pull request? This PR adds a shared DataStreamReader to sql. ### Why are the changes needed? We are creating a unified Scala interface for sql. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48213 from hvanhovell/SPARK-49416. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/SparkSession.scala | 10 +- .../sql/streaming/DataStreamReader.scala | 295 ++++------------ .../CheckConnectJvmClientCompatibility.scala | 8 +- project/MimaExcludes.scala | 1 + .../spark/sql/api/DataStreamReader.scala | 297 ++++++++++++++++ .../apache/spark/sql/api/SparkSession.scala | 11 + .../org/apache/spark/sql/SparkSession.scala | 10 +- .../sql/streaming/DataStreamReader.scala | 325 ++++-------------- 8 files changed, 438 insertions(+), 519 deletions(-) create mode 100644 sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamReader.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 5313369a2c987..1b41566ca1d1d 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -209,15 +209,7 @@ class SparkSession private[sql] ( /** @inheritdoc */ def read: DataFrameReader = new DataFrameReader(this) - /** - * Returns a `DataStreamReader` that can be used to read streaming data in as a `DataFrame`. - * {{{ - * sparkSession.readStream.parquet("/path/to/directory/of/parquet/files") - * sparkSession.readStream.schema(schema).json("/path/to/directory/of/json/files") - * }}} - * - * @since 3.5.0 - */ + /** @inheritdoc */ def readStream: DataStreamReader = new DataStreamReader(this) lazy val streams: StreamingQueryManager = new StreamingQueryManager(this) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala index 789425c9daea1..2ff34a6343644 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala @@ -21,11 +21,9 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Evolving import org.apache.spark.connect.proto.Read.DataSource -import org.apache.spark.internal.Logging -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.Dataset -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.StringEncoder +import org.apache.spark.sql.{api, DataFrame, Dataset, SparkSession} +import org.apache.spark.sql.connect.ConnectConversions._ +import org.apache.spark.sql.errors.DataTypeErrors import org.apache.spark.sql.types.StructType /** @@ -35,101 +33,49 @@ import org.apache.spark.sql.types.StructType * @since 3.5.0 */ @Evolving -final class DataStreamReader private[sql] (sparkSession: SparkSession) extends Logging { +final class DataStreamReader private[sql] (sparkSession: SparkSession) + extends api.DataStreamReader { - /** - * Specifies the input data source format. - * - * @since 3.5.0 - */ - def format(source: String): DataStreamReader = { + private val sourceBuilder = DataSource.newBuilder() + + /** @inheritdoc */ + def format(source: String): this.type = { sourceBuilder.setFormat(source) this } - /** - * Specifies the input schema. Some data sources (e.g. JSON) can infer the input schema - * automatically from data. By specifying the schema here, the underlying data source can skip - * the schema inference step, and thus speed up data loading. - * - * @since 3.5.0 - */ - def schema(schema: StructType): DataStreamReader = { + /** @inheritdoc */ + def schema(schema: StructType): this.type = { if (schema != null) { sourceBuilder.setSchema(schema.json) // Use json. DDL does not retail all the attributes. } this } - /** - * Specifies the schema by using the input DDL-formatted string. Some data sources (e.g. JSON) - * can infer the input schema automatically from data. By specifying the schema here, the - * underlying data source can skip the schema inference step, and thus speed up data loading. - * - * @since 3.5.0 - */ - def schema(schemaString: String): DataStreamReader = { + /** @inheritdoc */ + override def schema(schemaString: String): this.type = { sourceBuilder.setSchema(schemaString) this } - /** - * Adds an input option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: String): DataStreamReader = { + /** @inheritdoc */ + def option(key: String, value: String): this.type = { sourceBuilder.putOptions(key, value) this } - /** - * Adds an input option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Boolean): DataStreamReader = option(key, value.toString) - - /** - * Adds an input option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Long): DataStreamReader = option(key, value.toString) - - /** - * Adds an input option for the underlying data source. - * - * @since 3.5.0 - */ - def option(key: String, value: Double): DataStreamReader = option(key, value.toString) - - /** - * (Scala-specific) Adds input options for the underlying data source. - * - * @since 3.5.0 - */ - def options(options: scala.collection.Map[String, String]): DataStreamReader = { + /** @inheritdoc */ + def options(options: scala.collection.Map[String, String]): this.type = { this.options(options.asJava) - this } - /** - * (Java-specific) Adds input options for the underlying data source. - * - * @since 3.5.0 - */ - def options(options: java.util.Map[String, String]): DataStreamReader = { + /** @inheritdoc */ + override def options(options: java.util.Map[String, String]): this.type = { sourceBuilder.putAllOptions(options) this } - /** - * Loads input data stream in as a `DataFrame`, for data streams that don't require a path (e.g. - * external key-value stores). - * - * @since 3.5.0 - */ + /** @inheritdoc */ def load(): DataFrame = { sparkSession.newDataFrame { relationBuilder => relationBuilder.getReadBuilder @@ -138,120 +84,14 @@ final class DataStreamReader private[sql] (sparkSession: SparkSession) extends L } } - /** - * Loads input in as a `DataFrame`, for data streams that read from some path. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def load(path: String): DataFrame = { sourceBuilder.clearPaths() sourceBuilder.addPaths(path) load() } - /** - * Loads a JSON file stream and returns the results as a `DataFrame`. - * - * JSON Lines (newline-delimited JSON) is supported by - * default. For JSON (one record per file), set the `multiLine` option to true. - * - * This function goes through the input once to determine the input schema. If you know the - * schema in advance, use the version that specifies the schema to avoid the extra scan. - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * You can find the JSON-specific options for reading JSON file stream in - * Data Source Option in the version you use. - * - * @since 3.5.0 - */ - def json(path: String): DataFrame = { - format("json").load(path) - } - - /** - * Loads a CSV file stream and returns the result as a `DataFrame`. - * - * This function will go through the input once to determine the input schema if `inferSchema` - * is enabled. To avoid going through the entire data once, disable `inferSchema` option or - * specify the schema explicitly using `schema`. - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * You can find the CSV-specific options for reading CSV file stream in - * Data Source Option in the version you use. - * - * @since 3.5.0 - */ - def csv(path: String): DataFrame = format("csv").load(path) - - /** - * Loads a XML file stream and returns the result as a `DataFrame`. - * - * This function will go through the input once to determine the input schema if `inferSchema` - * is enabled. To avoid going through the entire data once, disable `inferSchema` option or - * specify the schema explicitly using `schema`. - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * You can find the XML-specific options for reading XML file stream in - * Data Source Option in the version you use. - * - * @since 4.0.0 - */ - def xml(path: String): DataFrame = format("xml").load(path) - - /** - * Loads a ORC file stream, returning the result as a `DataFrame`. - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * ORC-specific option(s) for reading ORC file stream can be found in Data - * Source Option in the version you use. - * - * @since 3.5.0 - */ - def orc(path: String): DataFrame = format("orc").load(path) - - /** - * Loads a Parquet file stream, returning the result as a `DataFrame`. - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * Parquet-specific option(s) for reading Parquet file stream can be found in Data - * Source Option in the version you use. - * - * @since 3.5.0 - */ - def parquet(path: String): DataFrame = format("parquet").load(path) - - /** - * Define a Streaming DataFrame on a Table. The DataSource corresponding to the table should - * support streaming mode. - * @param tableName - * The name of the table - * @since 3.5.0 - */ + /** @inheritdoc */ def table(tableName: String): DataFrame = { require(tableName != null, "The table name can't be null") sparkSession.newDataFrame { builder => @@ -263,59 +103,44 @@ final class DataStreamReader private[sql] (sparkSession: SparkSession) extends L } } - /** - * Loads text files and returns a `DataFrame` whose schema starts with a string column named - * "value", and followed by partitioned columns if there are any. The text files must be encoded - * as UTF-8. - * - * By default, each line in the text files is a new row in the resulting DataFrame. For example: - * {{{ - * // Scala: - * spark.readStream.text("/path/to/directory/") - * - * // Java: - * spark.readStream().text("/path/to/directory/") - * }}} - * - * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): - * sets the maximum number of new files to be considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to - * be considered in every trigger.
    - * - * You can find the text-specific options for reading text files in - * Data Source Option in the version you use. - * - * @since 3.5.0 - */ - def text(path: String): DataFrame = format("text").load(path) - - /** - * Loads text file(s) and returns a `Dataset` of String. The underlying schema of the Dataset - * contains a single string column named "value". The text files must be encoded as UTF-8. - * - * If the directory structure of the text files contains partitioning information, those are - * ignored in the resulting Dataset. To include partitioning information as columns, use `text`. - * - * By default, each line in the text file is a new element in the resulting Dataset. For - * example: - * {{{ - * // Scala: - * spark.readStream.textFile("/path/to/spark/README.md") - * - * // Java: - * spark.readStream().textFile("/path/to/spark/README.md") - * }}} - * - * You can set the text-specific options as specified in `DataStreamReader.text`. - * - * @param path - * input path - * @since 3.5.0 - */ - def textFile(path: String): Dataset[String] = { - text(path).select("value").as[String](StringEncoder) + override protected def assertNoSpecifiedSchema(operation: String): Unit = { + if (sourceBuilder.hasSchema) { + throw DataTypeErrors.userSpecifiedSchemaUnsupportedError(operation) + } } - private val sourceBuilder = DataSource.newBuilder() + /////////////////////////////////////////////////////////////////////////////////////// + // Covariant overrides. + /////////////////////////////////////////////////////////////////////////////////////// + + /** @inheritdoc */ + override def option(key: String, value: Boolean): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Long): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Double): this.type = super.option(key, value) + + /** @inheritdoc */ + override def json(path: String): DataFrame = super.json(path) + + /** @inheritdoc */ + override def csv(path: String): DataFrame = super.csv(path) + + /** @inheritdoc */ + override def xml(path: String): DataFrame = super.xml(path) + + /** @inheritdoc */ + override def orc(path: String): DataFrame = super.orc(path) + + /** @inheritdoc */ + override def parquet(path: String): DataFrame = super.parquet(path) + + /** @inheritdoc */ + override def text(path: String): DataFrame = super.text(path) + + /** @inheritdoc */ + override def textFile(path: String): Dataset[String] = super.textFile(path) + } diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala index 16f6983efb187..c8776af18a14a 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala @@ -304,7 +304,13 @@ object CheckConnectJvmClientCompatibility { ProblemFilters.exclude[DirectMissingMethodProblem]( "org.apache.spark.sql.DataFrameReader.validateJsonSchema"), ProblemFilters.exclude[DirectMissingMethodProblem]( - "org.apache.spark.sql.DataFrameReader.validateXmlSchema")) + "org.apache.spark.sql.DataFrameReader.validateXmlSchema"), + + // Protected DataStreamReader methods... + ProblemFilters.exclude[DirectMissingMethodProblem]( + "org.apache.spark.sql.streaming.DataStreamReader.validateJsonSchema"), + ProblemFilters.exclude[DirectMissingMethodProblem]( + "org.apache.spark.sql.streaming.DataStreamReader.validateXmlSchema")) checkMiMaCompatibility(clientJar, sqlJar, includedRules, excludeRules) } diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 9a89ebb4797c9..0bd0121e6e141 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -179,6 +179,7 @@ object MimaExcludes { // SPARK-49282: Shared SparkSessionBuilder ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SparkSession$Builder"), ) ++ loggingExcludes("org.apache.spark.sql.DataFrameReader") ++ + loggingExcludes("org.apache.spark.sql.streaming.DataStreamReader") ++ loggingExcludes("org.apache.spark.sql.SparkSession#Builder") // Default exclude rules diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamReader.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamReader.scala new file mode 100644 index 0000000000000..219ecb77d4033 --- /dev/null +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/DataStreamReader.scala @@ -0,0 +1,297 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.api + +import scala.jdk.CollectionConverters._ + +import _root_.java + +import org.apache.spark.annotation.Evolving +import org.apache.spark.sql.{Encoders, Row} +import org.apache.spark.sql.types.StructType + +/** + * Interface used to load a streaming `Dataset` from external storage systems (e.g. file systems, + * key-value stores, etc). Use `SparkSession.readStream` to access this. + * + * @since 2.0.0 + */ +@Evolving +abstract class DataStreamReader { + + /** + * Specifies the input data source format. + * + * @since 2.0.0 + */ + def format(source: String): this.type + + /** + * Specifies the input schema. Some data sources (e.g. JSON) can infer the input schema + * automatically from data. By specifying the schema here, the underlying data source can skip + * the schema inference step, and thus speed up data loading. + * + * @since 2.0.0 + */ + def schema(schema: StructType): this.type + + /** + * Specifies the schema by using the input DDL-formatted string. Some data sources (e.g. JSON) + * can infer the input schema automatically from data. By specifying the schema here, the + * underlying data source can skip the schema inference step, and thus speed up data loading. + * + * @since 2.3.0 + */ + def schema(schemaString: String): this.type = { + schema(StructType.fromDDL(schemaString)) + } + + /** + * Adds an input option for the underlying data source. + * + * @since 2.0.0 + */ + def option(key: String, value: String): this.type + + /** + * Adds an input option for the underlying data source. + * + * @since 2.0.0 + */ + def option(key: String, value: Boolean): this.type = option(key, value.toString) + + /** + * Adds an input option for the underlying data source. + * + * @since 2.0.0 + */ + def option(key: String, value: Long): this.type = option(key, value.toString) + + /** + * Adds an input option for the underlying data source. + * + * @since 2.0.0 + */ + def option(key: String, value: Double): this.type = option(key, value.toString) + + /** + * (Scala-specific) Adds input options for the underlying data source. + * + * @since 2.0.0 + */ + def options(options: scala.collection.Map[String, String]): this.type + + /** + * (Java-specific) Adds input options for the underlying data source. + * + * @since 2.0.0 + */ + def options(options: java.util.Map[String, String]): this.type = { + this.options(options.asScala) + this + } + + /** + * Loads input data stream in as a `DataFrame`, for data streams that don't require a path (e.g. + * external key-value stores). + * + * @since 2.0.0 + */ + def load(): Dataset[Row] + + /** + * Loads input in as a `DataFrame`, for data streams that read from some path. + * + * @since 2.0.0 + */ + def load(path: String): Dataset[Row] + + /** + * Loads a JSON file stream and returns the results as a `DataFrame`. + * + * JSON Lines (newline-delimited JSON) is supported by + * default. For JSON (one record per file), set the `multiLine` option to true. + * + * This function goes through the input once to determine the input schema. If you know the + * schema in advance, use the version that specifies the schema to avoid the extra scan. + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * You can find the JSON-specific options for reading JSON file stream in + * Data Source Option in the version you use. + * + * @since 2.0.0 + */ + def json(path: String): Dataset[Row] = { + validateJsonSchema() + format("json").load(path) + } + + /** + * Loads a CSV file stream and returns the result as a `DataFrame`. + * + * This function will go through the input once to determine the input schema if `inferSchema` + * is enabled. To avoid going through the entire data once, disable `inferSchema` option or + * specify the schema explicitly using `schema`. + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * You can find the CSV-specific options for reading CSV file stream in + * Data Source Option in the version you use. + * + * @since 2.0.0 + */ + def csv(path: String): Dataset[Row] = format("csv").load(path) + + /** + * Loads a XML file stream and returns the result as a `DataFrame`. + * + * This function will go through the input once to determine the input schema if `inferSchema` + * is enabled. To avoid going through the entire data once, disable `inferSchema` option or + * specify the schema explicitly using `schema`. + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * You can find the XML-specific options for reading XML file stream in + * Data Source Option in the version you use. + * + * @since 4.0.0 + */ + def xml(path: String): Dataset[Row] = { + validateXmlSchema() + format("xml").load(path) + } + + /** + * Loads a ORC file stream, returning the result as a `DataFrame`. + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * ORC-specific option(s) for reading ORC file stream can be found in Data + * Source Option in the version you use. + * + * @since 2.3.0 + */ + def orc(path: String): Dataset[Row] = { + format("orc").load(path) + } + + /** + * Loads a Parquet file stream, returning the result as a `DataFrame`. + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * Parquet-specific option(s) for reading Parquet file stream can be found in Data + * Source Option in the version you use. + * + * @since 2.0.0 + */ + def parquet(path: String): Dataset[Row] = { + format("parquet").load(path) + } + + /** + * Define a Streaming DataFrame on a Table. The DataSource corresponding to the table should + * support streaming mode. + * @param tableName + * The name of the table + * @since 3.1.0 + */ + def table(tableName: String): Dataset[Row] + + /** + * Loads text files and returns a `DataFrame` whose schema starts with a string column named + * "value", and followed by partitioned columns if there are any. The text files must be encoded + * as UTF-8. + * + * By default, each line in the text files is a new row in the resulting DataFrame. For example: + * {{{ + * // Scala: + * spark.readStream.text("/path/to/directory/") + * + * // Java: + * spark.readStream().text("/path/to/directory/") + * }}} + * + * You can set the following option(s):
    • `maxFilesPerTrigger` (default: no max limit): + * sets the maximum number of new files to be considered in every trigger.
      • + *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files to + * be considered in every trigger.
    + * + * You can find the text-specific options for reading text files in + * Data Source Option in the version you use. + * + * @since 2.0.0 + */ + def text(path: String): Dataset[Row] = format("text").load(path) + + /** + * Loads text file(s) and returns a `Dataset` of String. The underlying schema of the Dataset + * contains a single string column named "value". The text files must be encoded as UTF-8. + * + * If the directory structure of the text files contains partitioning information, those are + * ignored in the resulting Dataset. To include partitioning information as columns, use `text`. + * + * By default, each line in the text file is a new element in the resulting Dataset. For + * example: + * {{{ + * // Scala: + * spark.readStream.textFile("/path/to/spark/README.md") + * + * // Java: + * spark.readStream().textFile("/path/to/spark/README.md") + * }}} + * + * You can set the text-specific options as specified in `DataStreamReader.text`. + * + * @param path + * input path + * @since 2.1.0 + */ + def textFile(path: String): Dataset[String] = { + assertNoSpecifiedSchema("textFile") + text(path).select("value").as(Encoders.STRING) + } + + protected def assertNoSpecifiedSchema(operation: String): Unit + + protected def validateJsonSchema(): Unit = () + + protected def validateXmlSchema(): Unit = () + +} diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 2295c153cd51c..0f73a94c3c4a4 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -506,6 +506,17 @@ abstract class SparkSession extends Serializable with Closeable { */ def read: DataFrameReader + /** + * Returns a `DataStreamReader` that can be used to read streaming data in as a `DataFrame`. + * {{{ + * sparkSession.readStream.parquet("/path/to/directory/of/parquet/files") + * sparkSession.readStream.schema(schema).json("/path/to/directory/of/json/files") + * }}} + * + * @since 2.0.0 + */ + def readStream: DataStreamReader + /** * (Scala-specific) Implicit methods available in Scala for converting common Scala objects into * `DataFrame`s. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index fe139d629eb24..983cc24718fd2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -739,15 +739,7 @@ class SparkSession private( /** @inheritdoc */ def read: DataFrameReader = new DataFrameReader(self) - /** - * Returns a `DataStreamReader` that can be used to read streaming data in as a `DataFrame`. - * {{{ - * sparkSession.readStream.parquet("/path/to/directory/of/parquet/files") - * sparkSession.readStream.schema(schema).json("/path/to/directory/of/json/files") - * }}} - * - * @since 2.0.0 - */ + /** @inheritdoc */ def readStream: DataStreamReader = new DataStreamReader(self) // scalastyle:off diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala index 24d769fc8fc87..f42d8b667ab12 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala @@ -22,12 +22,12 @@ import java.util.Locale import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Evolving -import org.apache.spark.internal.Logging -import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} +import org.apache.spark.sql.{api, DataFrame, Dataset, SparkSession} import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, CharVarcharUtils} +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.connector.catalog.{SupportsRead, TableProvider} import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.errors.QueryCompilationErrors @@ -49,25 +49,15 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap * @since 2.0.0 */ @Evolving -final class DataStreamReader private[sql](sparkSession: SparkSession) extends Logging { - /** - * Specifies the input data source format. - * - * @since 2.0.0 - */ - def format(source: String): DataStreamReader = { +final class DataStreamReader private[sql](sparkSession: SparkSession) extends api.DataStreamReader { + /** @inheritdoc */ + def format(source: String): this.type = { this.source = source this } - /** - * Specifies the input schema. Some data sources (e.g. JSON) can infer the input schema - * automatically from data. By specifying the schema here, the underlying data source can - * skip the schema inference step, and thus speed up data loading. - * - * @since 2.0.0 - */ - def schema(schema: StructType): DataStreamReader = { + /** @inheritdoc */ + def schema(schema: StructType): this.type = { if (schema != null) { val replaced = CharVarcharUtils.failIfHasCharVarchar(schema).asInstanceOf[StructType] this.userSpecifiedSchema = Option(replaced) @@ -75,75 +65,19 @@ final class DataStreamReader private[sql](sparkSession: SparkSession) extends Lo this } - /** - * Specifies the schema by using the input DDL-formatted string. Some data sources (e.g. JSON) can - * infer the input schema automatically from data. By specifying the schema here, the underlying - * data source can skip the schema inference step, and thus speed up data loading. - * - * @since 2.3.0 - */ - def schema(schemaString: String): DataStreamReader = { - schema(StructType.fromDDL(schemaString)) - } - - /** - * Adds an input option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: String): DataStreamReader = { + /** @inheritdoc */ + def option(key: String, value: String): this.type = { this.extraOptions += (key -> value) this } - /** - * Adds an input option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Boolean): DataStreamReader = option(key, value.toString) - - /** - * Adds an input option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Long): DataStreamReader = option(key, value.toString) - - /** - * Adds an input option for the underlying data source. - * - * @since 2.0.0 - */ - def option(key: String, value: Double): DataStreamReader = option(key, value.toString) - - /** - * (Scala-specific) Adds input options for the underlying data source. - * - * @since 2.0.0 - */ - def options(options: scala.collection.Map[String, String]): DataStreamReader = { + /** @inheritdoc */ + def options(options: scala.collection.Map[String, String]): this.type = { this.extraOptions ++= options this } - /** - * (Java-specific) Adds input options for the underlying data source. - * - * @since 2.0.0 - */ - def options(options: java.util.Map[String, String]): DataStreamReader = { - this.options(options.asScala) - this - } - - - /** - * Loads input data stream in as a `DataFrame`, for data streams that don't require a path - * (e.g. external key-value stores). - * - * @since 2.0.0 - */ + /** @inheritdoc */ def load(): DataFrame = loadInternal(None) private def loadInternal(path: Option[String]): DataFrame = { @@ -205,11 +139,7 @@ final class DataStreamReader private[sql](sparkSession: SparkSession) extends Lo } } - /** - * Loads input in as a `DataFrame`, for data streams that read from some path. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def load(path: String): DataFrame = { if (!sparkSession.sessionState.conf.legacyPathOptionBehavior && extraOptions.contains("path")) { @@ -218,133 +148,7 @@ final class DataStreamReader private[sql](sparkSession: SparkSession) extends Lo loadInternal(Some(path)) } - /** - * Loads a JSON file stream and returns the results as a `DataFrame`. - * - * JSON Lines (newline-delimited JSON) is supported by - * default. For JSON (one record per file), set the `multiLine` option to true. - * - * This function goes through the input once to determine the input schema. If you know the - * schema in advance, use the version that specifies the schema to avoid the extra scan. - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * You can find the JSON-specific options for reading JSON file stream in - * - * Data Source Option in the version you use. - * - * @since 2.0.0 - */ - def json(path: String): DataFrame = { - userSpecifiedSchema.foreach(checkJsonSchema) - format("json").load(path) - } - - /** - * Loads a CSV file stream and returns the result as a `DataFrame`. - * - * This function will go through the input once to determine the input schema if `inferSchema` - * is enabled. To avoid going through the entire data once, disable `inferSchema` option or - * specify the schema explicitly using `schema`. - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * You can find the CSV-specific options for reading CSV file stream in - * - * Data Source Option in the version you use. - * - * @since 2.0.0 - */ - def csv(path: String): DataFrame = format("csv").load(path) - - /** - * Loads a XML file stream and returns the result as a `DataFrame`. - * - * This function will go through the input once to determine the input schema if `inferSchema` - * is enabled. To avoid going through the entire data once, disable `inferSchema` option or - * specify the schema explicitly using `schema`. - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * You can find the XML-specific options for reading XML file stream in - * - * Data Source Option in the version you use. - * - * @since 4.0.0 - */ - def xml(path: String): DataFrame = { - userSpecifiedSchema.foreach(checkXmlSchema) - format("xml").load(path) - } - - /** - * Loads a ORC file stream, returning the result as a `DataFrame`. - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * ORC-specific option(s) for reading ORC file stream can be found in - * - * Data Source Option in the version you use. - * - * @since 2.3.0 - */ - def orc(path: String): DataFrame = { - format("orc").load(path) - } - - /** - * Loads a Parquet file stream, returning the result as a `DataFrame`. - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * Parquet-specific option(s) for reading Parquet file stream can be found in - * - * Data Source Option in the version you use. - * - * @since 2.0.0 - */ - def parquet(path: String): DataFrame = { - format("parquet").load(path) - } - - /** - * Define a Streaming DataFrame on a Table. The DataSource corresponding to the table should - * support streaming mode. - * @param tableName The name of the table - * @since 3.1.0 - */ + /** @inheritdoc */ def table(tableName: String): DataFrame = { require(tableName != null, "The table name can't be null") val identifier = sparkSession.sessionState.sqlParser.parseMultipartIdentifier(tableName) @@ -356,65 +160,56 @@ final class DataStreamReader private[sql](sparkSession: SparkSession) extends Lo isStreaming = true)) } - /** - * Loads text files and returns a `DataFrame` whose schema starts with a string column named - * "value", and followed by partitioned columns if there are any. - * The text files must be encoded as UTF-8. - * - * By default, each line in the text files is a new row in the resulting DataFrame. For example: - * {{{ - * // Scala: - * spark.readStream.text("/path/to/directory/") - * - * // Java: - * spark.readStream().text("/path/to/directory/") - * }}} - * - * You can set the following option(s): - *
      - *
    • `maxFilesPerTrigger` (default: no max limit): sets the maximum number of new files to be - * considered in every trigger.
      • - *
    • `maxBytesPerTrigger` (default: no max limit): sets the maximum total size of new files - * to be considered in every trigger.
      • - *
    - * - * You can find the text-specific options for reading text files in - * - * Data Source Option in the version you use. - * - * @since 2.0.0 - */ - def text(path: String): DataFrame = format("text").load(path) - - /** - * Loads text file(s) and returns a `Dataset` of String. The underlying schema of the Dataset - * contains a single string column named "value". - * The text files must be encoded as UTF-8. - * - * If the directory structure of the text files contains partitioning information, those are - * ignored in the resulting Dataset. To include partitioning information as columns, use `text`. - * - * By default, each line in the text file is a new element in the resulting Dataset. For example: - * {{{ - * // Scala: - * spark.readStream.textFile("/path/to/spark/README.md") - * - * // Java: - * spark.readStream().textFile("/path/to/spark/README.md") - * }}} - * - * You can set the text-specific options as specified in `DataStreamReader.text`. - * - * @param path input path - * @since 2.1.0 - */ - def textFile(path: String): Dataset[String] = { + override protected def assertNoSpecifiedSchema(operation: String): Unit = { if (userSpecifiedSchema.nonEmpty) { - throw QueryCompilationErrors.userSpecifiedSchemaUnsupportedError("textFile") + throw QueryCompilationErrors.userSpecifiedSchemaUnsupportedError(operation) } - text(path).select("value").as[String](sparkSession.implicits.newStringEncoder) } + override protected def validateJsonSchema(): Unit = userSpecifiedSchema.foreach(checkJsonSchema) + + override protected def validateXmlSchema(): Unit = userSpecifiedSchema.foreach(checkXmlSchema) + + /////////////////////////////////////////////////////////////////////////////////////// + // Covariant overrides. + /////////////////////////////////////////////////////////////////////////////////////// + + /** @inheritdoc */ + override def schema(schemaString: String): this.type = super.schema(schemaString) + + /** @inheritdoc */ + override def option(key: String, value: Boolean): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Long): this.type = super.option(key, value) + + /** @inheritdoc */ + override def option(key: String, value: Double): this.type = super.option(key, value) + + /** @inheritdoc */ + override def options(options: java.util.Map[String, String]): this.type = super.options(options) + + /** @inheritdoc */ + override def json(path: String): DataFrame = super.json(path) + + /** @inheritdoc */ + override def csv(path: String): DataFrame = super.csv(path) + + /** @inheritdoc */ + override def xml(path: String): DataFrame = super.xml(path) + + /** @inheritdoc */ + override def orc(path: String): DataFrame = super.orc(path) + + /** @inheritdoc */ + override def parquet(path: String): DataFrame = super.parquet(path) + + /** @inheritdoc */ + override def text(path: String): DataFrame = super.text(path) + + /** @inheritdoc */ + override def textFile(path: String): Dataset[String] = super.textFile(path) + /////////////////////////////////////////////////////////////////////////////////////// // Builder pattern config options /////////////////////////////////////////////////////////////////////////////////////// From 7b20e5841a856cd0d81821e330b3ec33098bb9be Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Thu, 26 Sep 2024 09:28:16 -0400 Subject: [PATCH 146/189] [SPARK-49286][CONNECT][SQL] Move Avro/Protobuf functions to sql/api ### What changes were proposed in this pull request? This PR moves avro and protobuf functions to sql/api. ### Why are the changes needed? We are creating a unified Scala SQL interface. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48258 from hvanhovell/SPARK-49286. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/avro/functions.scala | 93 ----- .../apache/spark/sql/protobuf/functions.scala | 324 ------------------ project/MimaExcludes.scala | 6 + .../org/apache/spark/sql/avro/functions.scala | 8 +- .../apache/spark/sql/protobuf/functions.scala | 82 +++-- 5 files changed, 50 insertions(+), 463 deletions(-) delete mode 100755 connector/avro/src/main/scala/org/apache/spark/sql/avro/functions.scala delete mode 100644 connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/functions.scala rename {connector/connect/client/jvm => sql/api}/src/main/scala/org/apache/spark/sql/avro/functions.scala (97%) rename {connector/connect/client/jvm => sql/api}/src/main/scala/org/apache/spark/sql/protobuf/functions.scala (90%) diff --git a/connector/avro/src/main/scala/org/apache/spark/sql/avro/functions.scala b/connector/avro/src/main/scala/org/apache/spark/sql/avro/functions.scala deleted file mode 100755 index 828a609a10e9c..0000000000000 --- a/connector/avro/src/main/scala/org/apache/spark/sql/avro/functions.scala +++ /dev/null @@ -1,93 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.spark.sql.avro - -import scala.jdk.CollectionConverters._ - -import org.apache.spark.annotation.Experimental -import org.apache.spark.sql.Column -import org.apache.spark.sql.internal.ExpressionUtils.{column, expression} - - -// scalastyle:off: object.name -object functions { -// scalastyle:on: object.name - - /** - * Converts a binary column of avro format into its corresponding catalyst value. The specified - * schema must match the read data, otherwise the behavior is undefined: it may fail or return - * arbitrary result. - * - * @param data the binary column. - * @param jsonFormatSchema the avro schema in JSON string format. - * - * @since 3.0.0 - */ - @Experimental - def from_avro( - data: Column, - jsonFormatSchema: String): Column = { - AvroDataToCatalyst(data, jsonFormatSchema, Map.empty) - } - - /** - * Converts a binary column of Avro format into its corresponding catalyst value. - * The specified schema must match actual schema of the read data, otherwise the behavior - * is undefined: it may fail or return arbitrary result. - * To deserialize the data with a compatible and evolved schema, the expected Avro schema can be - * set via the option avroSchema. - * - * @param data the binary column. - * @param jsonFormatSchema the avro schema in JSON string format. - * @param options options to control how the Avro record is parsed. - * - * @since 3.0.0 - */ - @Experimental - def from_avro( - data: Column, - jsonFormatSchema: String, - options: java.util.Map[String, String]): Column = { - AvroDataToCatalyst(data, jsonFormatSchema, options.asScala.toMap) - } - - /** - * Converts a column into binary of avro format. - * - * @param data the data column. - * - * @since 3.0.0 - */ - @Experimental - def to_avro(data: Column): Column = { - CatalystDataToAvro(data, None) - } - - /** - * Converts a column into binary of avro format. - * - * @param data the data column. - * @param jsonFormatSchema user-specified output avro schema in JSON string format. - * - * @since 3.0.0 - */ - @Experimental - def to_avro(data: Column, jsonFormatSchema: String): Column = { - CatalystDataToAvro(data, Some(jsonFormatSchema)) - } -} diff --git a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/functions.scala b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/functions.scala deleted file mode 100644 index 3b0def8fc73f7..0000000000000 --- a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/functions.scala +++ /dev/null @@ -1,324 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.spark.sql.protobuf - -import scala.jdk.CollectionConverters._ - -import org.apache.spark.annotation.Experimental -import org.apache.spark.sql.Column -import org.apache.spark.sql.functions.lit -import org.apache.spark.sql.protobuf.utils.ProtobufUtils - -// scalastyle:off: object.name -object functions { -// scalastyle:on: object.name - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value. The - * Protobuf definition is provided through Protobuf descriptor file. - * - * @param data - * the binary column. - * @param messageName - * the protobuf message name to look for in descriptor file. - * @param descFilePath - * The Protobuf descriptor file. This file is usually created using `protoc` with - * `--descriptor_set_out` and `--include_imports` options. - * @param options - * @since 3.4.0 - */ - @Experimental - def from_protobuf( - data: Column, - messageName: String, - descFilePath: String, - options: java.util.Map[String, String]): Column = { - val descriptorFileContent = ProtobufUtils.readDescriptorFileContent(descFilePath) - from_protobuf(data, messageName, descriptorFileContent, options) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value.The - * Protobuf definition is provided through Protobuf `FileDescriptorSet`. - * - * @param data - * the binary column. - * @param messageName - * the protobuf MessageName to look for in the descriptor set. - * @param binaryFileDescriptorSet - * Serialized Protobuf descriptor (`FileDescriptorSet`). Typically contents of file created - * using `protoc` with `--descriptor_set_out` and `--include_imports` options. - * @param options - * @since 3.5.0 - */ - @Experimental - def from_protobuf( - data: Column, - messageName: String, - binaryFileDescriptorSet: Array[Byte], - options: java.util.Map[String, String]): Column = { - Column.fnWithOptions( - "from_protobuf", - options.asScala.iterator, - data, - lit(messageName), - lit(binaryFileDescriptorSet) - ) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value. The - * Protobuf definition is provided through Protobuf descriptor file. - * - * @param data - * the binary column. - * @param messageName - * the protobuf MessageName to look for in descriptor file. - * @param descFilePath - * The Protobuf descriptor file. This file is usually created using `protoc` with - * `--descriptor_set_out` and `--include_imports` options. - * @since 3.4.0 - */ - @Experimental - def from_protobuf(data: Column, messageName: String, descFilePath: String): Column = { - val fileContent = ProtobufUtils.readDescriptorFileContent(descFilePath) - from_protobuf(data, messageName, fileContent) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value.The - * Protobuf definition is provided through Protobuf `FileDescriptorSet`. - * - * @param data - * the binary column. - * @param messageName - * the protobuf MessageName to look for in the descriptor set. - * @param binaryFileDescriptorSet - * Serialized Protobuf descriptor (`FileDescriptorSet`). Typically contents of file created - * using `protoc` with `--descriptor_set_out` and `--include_imports` options. - * @since 3.5.0 - */ - @Experimental - def from_protobuf(data: Column, messageName: String, binaryFileDescriptorSet: Array[Byte]) - : Column = { - Column.fn( - "from_protobuf", - data, - lit(messageName), - lit(binaryFileDescriptorSet) - ) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value. - * `messageClassName` points to Protobuf Java class. The jar containing Java class should be - * shaded. Specifically, `com.google.protobuf.*` should be shaded to - * `org.sparkproject.spark_protobuf.protobuf.*`. - * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from - * Protobuf files. - * - * @param data - * the binary column. - * @param messageClassName - * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. - * The jar with these classes needs to be shaded as described above. - * @since 3.4.0 - */ - @Experimental - def from_protobuf(data: Column, messageClassName: String): Column = { - Column.fn( - "from_protobuf", - data, - lit(messageClassName) - ) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value. - * `messageClassName` points to Protobuf Java class. The jar containing Java class should be - * shaded. Specifically, `com.google.protobuf.*` should be shaded to - * `org.sparkproject.spark_protobuf.protobuf.*`. - * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from - * Protobuf files. - * - * @param data - * the binary column. - * @param messageClassName - * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. - * The jar with these classes needs to be shaded as described above. - * @param options - * @since 3.4.0 - */ - @Experimental - def from_protobuf( - data: Column, - messageClassName: String, - options: java.util.Map[String, String]): Column = { - Column.fnWithOptions( - "from_protobuf", - options.asScala.iterator, - data, - lit(messageClassName) - ) - } - - /** - * Converts a column into binary of protobuf format. The Protobuf definition is provided - * through Protobuf descriptor file. - * - * @param data - * the data column. - * @param messageName - * the protobuf MessageName to look for in descriptor file. - * @param descFilePath - * The Protobuf descriptor file. This file is usually created using `protoc` with - * `--descriptor_set_out` and `--include_imports` options. - * @since 3.4.0 - */ - @Experimental - def to_protobuf(data: Column, messageName: String, descFilePath: String): Column = { - to_protobuf(data, messageName, descFilePath, Map.empty[String, String].asJava) - } - - /** - * Converts a column into binary of protobuf format.The Protobuf definition is provided - * through Protobuf `FileDescriptorSet`. - * - * @param data - * the binary column. - * @param messageName - * the protobuf MessageName to look for in the descriptor set. - * @param binaryFileDescriptorSet - * Serialized Protobuf descriptor (`FileDescriptorSet`). Typically contents of file created - * using `protoc` with `--descriptor_set_out` and `--include_imports` options. - * - * @since 3.5.0 - */ - @Experimental - def to_protobuf(data: Column, messageName: String, binaryFileDescriptorSet: Array[Byte]) - : Column = { - Column.fn( - "to_protobuf", - data, - lit(messageName), - lit(binaryFileDescriptorSet) - ) - } - /** - * Converts a column into binary of protobuf format. The Protobuf definition is provided - * through Protobuf descriptor file. - * - * @param data - * the data column. - * @param messageName - * the protobuf MessageName to look for in descriptor file. - * @param descFilePath - * the protobuf descriptor file. - * @param options - * @since 3.4.0 - */ - @Experimental - def to_protobuf( - data: Column, - messageName: String, - descFilePath: String, - options: java.util.Map[String, String]): Column = { - val fileContent = ProtobufUtils.readDescriptorFileContent(descFilePath) - to_protobuf(data, messageName, fileContent, options) - } - - /** - * Converts a column into binary of protobuf format.The Protobuf definition is provided - * through Protobuf `FileDescriptorSet`. - * - * @param data - * the binary column. - * @param messageName - * the protobuf MessageName to look for in the descriptor set. - * @param binaryFileDescriptorSet - * Serialized Protobuf descriptor (`FileDescriptorSet`). Typically contents of file created - * using `protoc` with `--descriptor_set_out` and `--include_imports` options. - * @param options - * @since 3.5.0 - */ - @Experimental - def to_protobuf( - data: Column, - messageName: String, - binaryFileDescriptorSet: Array[Byte], - options: java.util.Map[String, String] - ): Column = { - Column.fnWithOptions( - "to_protobuf", - options.asScala.iterator, - data, - lit(messageName), - lit(binaryFileDescriptorSet) - ) - } - - /** - * Converts a column into binary of protobuf format. - * `messageClassName` points to Protobuf Java class. The jar containing Java class should be - * shaded. Specifically, `com.google.protobuf.*` should be shaded to - * `org.sparkproject.spark_protobuf.protobuf.*`. - * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from - * Protobuf files. - * - * @param data - * the data column. - * @param messageClassName - * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. - * The jar with these classes needs to be shaded as described above. - * @since 3.4.0 - */ - @Experimental - def to_protobuf(data: Column, messageClassName: String): Column = { - Column.fn( - "to_protobuf", - data, - lit(messageClassName) - ) - } - - /** - * Converts a column into binary of protobuf format. - * `messageClassName` points to Protobuf Java class. The jar containing Java class should be - * shaded. Specifically, `com.google.protobuf.*` should be shaded to - * `org.sparkproject.spark_protobuf.protobuf.*`. - * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from - * Protobuf files. - * - * @param data - * the data column. - * @param messageClassName - * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. - * The jar with these classes needs to be shaded as described above. - * @param options - * @since 3.4.0 - */ - @Experimental - def to_protobuf(data: Column, messageClassName: String, options: java.util.Map[String, String]) - : Column = { - Column.fnWithOptions( - "to_protobuf", - options.asScala.iterator, - data, - lit(messageClassName) - ) - } -} diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 0bd0121e6e141..41f547a43b698 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -178,6 +178,12 @@ object MimaExcludes { // SPARK-49282: Shared SparkSessionBuilder ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.SparkSession$Builder"), + + // SPARK-49286: Avro/Protobuf functions in sql/api + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.avro.functions"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.avro.functions$"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.protobuf.functions"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.protobuf.functions$"), ) ++ loggingExcludes("org.apache.spark.sql.DataFrameReader") ++ loggingExcludes("org.apache.spark.sql.streaming.DataStreamReader") ++ loggingExcludes("org.apache.spark.sql.SparkSession#Builder") diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/avro/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/avro/functions.scala similarity index 97% rename from connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/avro/functions.scala rename to sql/api/src/main/scala/org/apache/spark/sql/avro/functions.scala index e80bccfee4c9c..fffad557aca5e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/avro/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/avro/functions.scala @@ -37,7 +37,7 @@ object functions { * @param jsonFormatSchema * the avro schema in JSON string format. * - * @since 3.5.0 + * @since 3.0.0 */ @Experimental def from_avro(data: Column, jsonFormatSchema: String): Column = { @@ -57,7 +57,7 @@ object functions { * @param options * options to control how the Avro record is parsed. * - * @since 3.5.0 + * @since 3.0.0 */ @Experimental def from_avro( @@ -73,7 +73,7 @@ object functions { * @param data * the data column. * - * @since 3.5.0 + * @since 3.0.0 */ @Experimental def to_avro(data: Column): Column = { @@ -88,7 +88,7 @@ object functions { * @param jsonFormatSchema * user-specified output avro schema in JSON string format. * - * @since 3.5.0 + * @since 3.0.0 */ @Experimental def to_avro(data: Column, jsonFormatSchema: String): Column = { diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/protobuf/functions.scala similarity index 90% rename from connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala rename to sql/api/src/main/scala/org/apache/spark/sql/protobuf/functions.scala index 2c953fbd07b9e..ea9e3c429d65a 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/protobuf/functions.scala @@ -18,7 +18,6 @@ package org.apache.spark.sql.protobuf import java.io.FileNotFoundException import java.nio.file.{Files, NoSuchFileException, Paths} -import java.util.Collections import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal @@ -30,7 +29,7 @@ import org.apache.spark.sql.functions.lit // scalastyle:off: object.name object functions { - // scalastyle:on: object.name +// scalastyle:on: object.name /** * Converts a binary column of Protobuf format into its corresponding catalyst value. The @@ -44,7 +43,7 @@ object functions { * The Protobuf descriptor file. This file is usually created using `protoc` with * `--descriptor_set_out` and `--include_imports` options. * @param options - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def from_protobuf( @@ -52,8 +51,8 @@ object functions { messageName: String, descFilePath: String, options: java.util.Map[String, String]): Column = { - val binaryFileDescSet = readDescriptorFileContent(descFilePath) - from_protobuf(data, messageName, binaryFileDescSet, options) + val descriptorFileContent = readDescriptorFileContent(descFilePath) + from_protobuf(data, messageName, descriptorFileContent, options) } /** @@ -95,31 +94,12 @@ object functions { * @param descFilePath * The Protobuf descriptor file. This file is usually created using `protoc` with * `--descriptor_set_out` and `--include_imports` options. - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def from_protobuf(data: Column, messageName: String, descFilePath: String): Column = { - from_protobuf(data, messageName, descFilePath, emptyOptions) - } - - /** - * Converts a binary column of Protobuf format into its corresponding catalyst value. - * `messageClassName` points to Protobuf Java class. The jar containing Java class should be - * shaded. Specifically, `com.google.protobuf.*` should be shaded to - * `org.sparkproject.spark_protobuf.protobuf.*`. - * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from - * Protobuf files. - * - * @param data - * the binary column. - * @param messageClassName - * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. - * The jar with these classes needs to be shaded as described above. - * @since 3.5.0 - */ - @Experimental - def from_protobuf(data: Column, messageClassName: String): Column = { - Column.fn("from_protobuf", data, lit(messageClassName)) + val fileContent = readDescriptorFileContent(descFilePath) + from_protobuf(data, messageName, fileContent) } /** @@ -140,7 +120,27 @@ object functions { data: Column, messageName: String, binaryFileDescriptorSet: Array[Byte]): Column = { - from_protobuf(data, messageName, binaryFileDescriptorSet, emptyOptions) + Column.fn("from_protobuf", data, lit(messageName), lit(binaryFileDescriptorSet)) + } + + /** + * Converts a binary column of Protobuf format into its corresponding catalyst value. + * `messageClassName` points to Protobuf Java class. The jar containing Java class should be + * shaded. Specifically, `com.google.protobuf.*` should be shaded to + * `org.sparkproject.spark_protobuf.protobuf.*`. + * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from + * Protobuf files. + * + * @param data + * the binary column. + * @param messageClassName + * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. + * The jar with these classes needs to be shaded as described above. + * @since 3.4.0 + */ + @Experimental + def from_protobuf(data: Column, messageClassName: String): Column = { + Column.fn("from_protobuf", data, lit(messageClassName)) } /** @@ -157,7 +157,7 @@ object functions { * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. * The jar with these classes needs to be shaded as described above. * @param options - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def from_protobuf( @@ -178,11 +178,11 @@ object functions { * @param descFilePath * The Protobuf descriptor file. This file is usually created using `protoc` with * `--descriptor_set_out` and `--include_imports` options. - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def to_protobuf(data: Column, messageName: String, descFilePath: String): Column = { - to_protobuf(data, messageName, descFilePath, emptyOptions) + to_protobuf(data, messageName, descFilePath, Map.empty[String, String].asJava) } /** @@ -204,7 +204,7 @@ object functions { data: Column, messageName: String, binaryFileDescriptorSet: Array[Byte]): Column = { - to_protobuf(data, messageName, binaryFileDescriptorSet, emptyOptions) + Column.fn("to_protobuf", data, lit(messageName), lit(binaryFileDescriptorSet)) } /** @@ -216,10 +216,9 @@ object functions { * @param messageName * the protobuf MessageName to look for in descriptor file. * @param descFilePath - * The Protobuf descriptor file. This file is usually created using `protoc` with - * `--descriptor_set_out` and `--include_imports` options. + * the protobuf descriptor file. * @param options - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def to_protobuf( @@ -227,8 +226,8 @@ object functions { messageName: String, descFilePath: String, options: java.util.Map[String, String]): Column = { - val binaryFileDescriptorSet = readDescriptorFileContent(descFilePath) - to_protobuf(data, messageName, binaryFileDescriptorSet, options) + val fileContent = readDescriptorFileContent(descFilePath) + to_protobuf(data, messageName, fileContent, options) } /** @@ -271,7 +270,7 @@ object functions { * @param messageClassName * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. * The jar with these classes needs to be shaded as described above. - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def to_protobuf(data: Column, messageClassName: String): Column = { @@ -291,7 +290,7 @@ object functions { * The full name for Protobuf Java class. E.g. com.example.protos.ExampleEvent. * The jar with these classes needs to be shaded as described above. * @param options - * @since 3.5.0 + * @since 3.4.0 */ @Experimental def to_protobuf( @@ -301,8 +300,6 @@ object functions { Column.fnWithOptions("to_protobuf", options.asScala.iterator, data, lit(messageClassName)) } - private def emptyOptions: java.util.Map[String, String] = Collections.emptyMap[String, String]() - // This method is copied from org.apache.spark.sql.protobuf.util.ProtobufUtils private def readDescriptorFileContent(filePath: String): Array[Byte] = { try { @@ -312,7 +309,8 @@ object functions { throw CompilationErrors.cannotFindDescriptorFileError(filePath, ex) case ex: NoSuchFileException => throw CompilationErrors.cannotFindDescriptorFileError(filePath, ex) - case NonFatal(ex) => throw CompilationErrors.descriptorParseError(ex) + case NonFatal(ex) => + throw CompilationErrors.descriptorParseError(ex) } } } From 218051a566c78244573077a53d4be43ccc01311d Mon Sep 17 00:00:00 2001 From: panbingkun Date: Thu, 26 Sep 2024 22:30:33 +0800 Subject: [PATCH 147/189] [MINOR][SQL][TESTS] Use `formatString.format(value)` instead of `value.formatted(formatString)` ### What changes were proposed in this pull request? The pr aims to use `formatString.format(value)` instead of `value.formatted(formatString)` for eliminating Warning. ### Why are the changes needed? image ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48262 from panbingkun/minor_formatted. Authored-by: panbingkun Signed-off-by: yangjie01 --- .../columnar/compression/CompressionSchemeBenchmark.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/CompressionSchemeBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/CompressionSchemeBenchmark.scala index 05ae575305299..290cfd56b8bce 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/CompressionSchemeBenchmark.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/CompressionSchemeBenchmark.scala @@ -91,7 +91,7 @@ object CompressionSchemeBenchmark extends BenchmarkBase with AllCompressionSchem schemes.filter(_.supports(tpe)).foreach { scheme => val (compressFunc, compressionRatio, buf) = prepareEncodeInternal(count, tpe, scheme, input) - val label = s"${getFormattedClassName(scheme)}(${compressionRatio.formatted("%.3f")})" + val label = s"${getFormattedClassName(scheme)}(${"%.3f".format(compressionRatio)})" benchmark.addCase(label)({ i: Int => for (n <- 0L until iters) { From 87b5ffb220824449d943cf3c7fff3eb3682526fc Mon Sep 17 00:00:00 2001 From: panbingkun Date: Thu, 26 Sep 2024 07:37:22 -0700 Subject: [PATCH 148/189] [SPARK-49797][INFRA] Align the running OS image of `maven_test.yml` to `ubuntu-latest` ### What changes were proposed in this pull request? The pr aims to align the running OS image of `maven_test.yml` to `ubuntu-latest` (from `ubuntu-22.04` to `ubuntu-24.04`) ### Why are the changes needed? https://github.com/actions/runner-images/releases/tag/ubuntu24%2F20240922.1 image After https://github.com/actions/runner-images/issues/10636, `ubuntu-latest` has already pointed to `ubuntu-24.04` instead of `ubuntu-22.04`. image I have checked all tasks running on `Ubuntu OS` (except for the 2 related to `TPCDS`), and they are all using `ubuntu-latest`. Currently, only `maven_test.yml` is using `ubuntu-22.04`. Let's align it. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48263 from panbingkun/SPARK-49797. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- .github/workflows/maven_test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/maven_test.yml b/.github/workflows/maven_test.yml index 82b72bd7e91d2..dd089d665d6e3 100644 --- a/.github/workflows/maven_test.yml +++ b/.github/workflows/maven_test.yml @@ -40,7 +40,7 @@ on: description: OS to run this build. required: false type: string - default: ubuntu-22.04 + default: ubuntu-latest envs: description: Additional environment variables to set when running the tests. Should be in JSON format. required: false From 624eda5030eb3a4a426a1c225952af40dba30d1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vladan=20Vasi=C4=87?= Date: Thu, 26 Sep 2024 23:00:22 +0800 Subject: [PATCH 149/189] [SPARK-49444][SQL] Modified UnivocityParser to throw runtime exceptions caused by ArrayIndexOutOfBounds with more user-oriented messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? I propose to catch and rethrow runtime `ArrayIndexOutOfBounds` exceptions in the `UnivocityParser` class - `parse` method, but with more user-oriented messages. Instead of throwing exceptions in the original format, I propose to inform the users which csv record caused the error. ### Why are the changes needed? Proper informing of users' errors improves user experience. Instead of throwing `ArrayIndexOutOfBounds` exception without clear reason why it happened, proposed changes throw `SparkRuntimeException` with the message that includes original csv line which caused the error. ### Does this PR introduce _any_ user-facing change? This PR introduces a user-facing change which happens when `UnivocityParser` parses malformed csv line with from the input. More specifically, the change is reproduces in the test case within `UnivocityParserSuite` when user specifies `maxColumns` in parser options and parsed csv record has more columns. Instead of resulting in `ArrayIndexOutOfBounds` like mentioned in the HMR ticket, users now get `SparkRuntimeException` with message that contains the input line which caused the error. ### How was this patch tested? This patch was tested in `UnivocityParserSuite`. Test named "Array index out of bounds when parsing CSV with more columns than expected" covers this patch. Additionally, test for bad records in `UnivocityParser`'s `PERMISSIVE` mode is added to confirm that `BadRecordException` is being thrown properly. ### Was this patch authored or co-authored using generative AI tooling? No Closes #47906 from vladanvasi-db/vladanvasi-db/univocity-parser-index-out-of-bounds-handling. Authored-by: Vladan Vasić Signed-off-by: Wenchen Fan --- .../sql/catalyst/csv/UnivocityParser.scala | 19 ++++++++- .../catalyst/csv/UnivocityParserSuite.scala | 39 ++++++++++++++++++- .../test/resources/test-data/more-columns.csv | 1 + .../apache/spark/sql/CsvFunctionsSuite.scala | 5 ++- .../execution/datasources/csv/CSVSuite.scala | 34 ++++++++++++++++ 5 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 sql/core/src/test/resources/test-data/more-columns.csv diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/UnivocityParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/UnivocityParser.scala index ccc8f30a9a9c3..0fd0601803a6a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/UnivocityParser.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/UnivocityParser.scala @@ -21,9 +21,10 @@ import java.io.InputStream import scala.util.control.NonFatal +import com.univocity.parsers.common.TextParsingException import com.univocity.parsers.csv.CsvParser -import org.apache.spark.SparkUpgradeException +import org.apache.spark.{SparkRuntimeException, SparkUpgradeException} import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.{InternalRow, NoopFilters, OrderedFilters} import org.apache.spark.sql.catalyst.expressions.{ExprUtils, GenericInternalRow} @@ -294,6 +295,20 @@ class UnivocityParser( } } + private def parseLine(line: String): Array[String] = { + try { + tokenizer.parseLine(line) + } + catch { + case e: TextParsingException if e.getCause.isInstanceOf[ArrayIndexOutOfBoundsException] => + throw new SparkRuntimeException( + errorClass = "MALFORMED_CSV_RECORD", + messageParameters = Map("badRecord" -> line), + cause = e + ) + } + } + /** * Parses a single CSV string and turns it into either one resulting row or no row (if the * the record is malformed). @@ -306,7 +321,7 @@ class UnivocityParser( (_: String) => Some(InternalRow.empty) } else { // parse if the columnPruning is disabled or requiredSchema is nonEmpty - (input: String) => convert(tokenizer.parseLine(input)) + (input: String) => convert(parseLine(input)) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/csv/UnivocityParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/csv/UnivocityParserSuite.scala index 514b529ea8cc0..7974bf68bdd31 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/csv/UnivocityParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/csv/UnivocityParserSuite.scala @@ -23,12 +23,12 @@ import java.util.{Locale, TimeZone} import org.apache.commons.lang3.time.FastDateFormat -import org.apache.spark.{SparkFunSuite, SparkIllegalArgumentException} +import org.apache.spark.{SparkFunSuite, SparkIllegalArgumentException, SparkRuntimeException} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.plans.SQLHelper +import org.apache.spark.sql.catalyst.util.{BadRecordException, DateTimeUtils} import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.catalyst.util.DateTimeTestUtils._ -import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.sources.{EqualTo, Filter, StringStartsWith} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -323,6 +323,41 @@ class UnivocityParserSuite extends SparkFunSuite with SQLHelper { parameters = Map("fieldName" -> "`i`", "fields" -> "")) } + test("Bad records test in permissive mode") { + def checkBadRecord( + input: String = "1,a", + dataSchema: StructType = StructType.fromDDL("i INTEGER, s STRING, d DOUBLE"), + requiredSchema: StructType = StructType.fromDDL("i INTEGER, s STRING"), + options: Map[String, String] = Map("mode" -> "PERMISSIVE")): BadRecordException = { + val csvOptions = new CSVOptions(options, false, "UTC") + val parser = new UnivocityParser(dataSchema, requiredSchema, csvOptions, Seq()) + intercept[BadRecordException] { + parser.parse(input) + } + } + + // Bad record exception caused by conversion error + checkBadRecord(input = "1.5,a,10.3") + + // Bad record exception caused by insufficient number of columns + checkBadRecord(input = "2") + } + + test("Array index out of bounds when parsing CSV with more columns than expected") { + val input = "1,string,3.14,5,7" + val dataSchema: StructType = StructType.fromDDL("i INTEGER, a STRING") + val requiredSchema: StructType = StructType.fromDDL("i INTEGER, a STRING") + val options = new CSVOptions(Map("maxColumns" -> "2"), false, "UTC") + val filters = Seq() + val parser = new UnivocityParser(dataSchema, requiredSchema, options, filters) + checkError( + exception = intercept[SparkRuntimeException] { + parser.parse(input) + }, + condition = "MALFORMED_CSV_RECORD", + parameters = Map("badRecord" -> "1,string,3.14,5,7")) + } + test("SPARK-30960: parse date/timestamp string with legacy format") { def check(parser: UnivocityParser): Unit = { // The legacy format allows 1 or 2 chars for some fields. diff --git a/sql/core/src/test/resources/test-data/more-columns.csv b/sql/core/src/test/resources/test-data/more-columns.csv new file mode 100644 index 0000000000000..06db38f0a145a --- /dev/null +++ b/sql/core/src/test/resources/test-data/more-columns.csv @@ -0,0 +1 @@ +1,3.14,string,5,7 \ No newline at end of file diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CsvFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CsvFunctionsSuite.scala index 6589282fd3a51..e6907b8656482 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CsvFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CsvFunctionsSuite.scala @@ -24,7 +24,8 @@ import java.util.Locale import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkException, SparkUnsupportedOperationException, SparkUpgradeException} +import org.apache.spark.{SparkException, SparkRuntimeException, + SparkUnsupportedOperationException, SparkUpgradeException} import org.apache.spark.sql.errors.DataTypeErrors.toSQLType import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -234,7 +235,7 @@ class CsvFunctionsSuite extends QueryTest with SharedSparkSession { val schema = new StructType().add("str", StringType) val options = Map("maxCharsPerColumn" -> "2") - val exception = intercept[SparkException] { + val exception = intercept[SparkRuntimeException] { df.select(from_csv($"value", schema, options)).collect() }.getCause.getMessage diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala index e2d1d9b05c3c2..023f401516dc3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala @@ -85,6 +85,7 @@ abstract class CSVSuite private val badAfterGoodFile = "test-data/bad_after_good.csv" private val malformedRowFile = "test-data/malformedRow.csv" private val charFile = "test-data/char.csv" + private val moreColumnsFile = "test-data/more-columns.csv" /** Verifies data and schema. */ private def verifyCars( @@ -3439,6 +3440,39 @@ abstract class CSVSuite expected) } } + + test("SPARK-49444: CSV parsing failure with more than max columns") { + val schema = new StructType() + .add("intColumn", IntegerType, nullable = true) + .add("decimalColumn", DecimalType(10, 2), nullable = true) + + val fileReadException = intercept[SparkException] { + spark + .read + .schema(schema) + .option("header", "false") + .option("maxColumns", "2") + .csv(testFile(moreColumnsFile)) + .collect() + } + + checkErrorMatchPVals( + exception = fileReadException, + condition = "FAILED_READ_FILE.NO_HINT", + parameters = Map("path" -> s".*$moreColumnsFile")) + + val malformedCSVException = fileReadException.getCause.asInstanceOf[SparkRuntimeException] + + checkError( + exception = malformedCSVException, + condition = "MALFORMED_CSV_RECORD", + parameters = Map("badRecord" -> "1,3.14,string,5,7"), + sqlState = "KD000") + + assert(malformedCSVException.getCause.isInstanceOf[TextParsingException]) + val textParsingException = malformedCSVException.getCause.asInstanceOf[TextParsingException] + assert(textParsingException.getCause.isInstanceOf[ArrayIndexOutOfBoundsException]) + } } class CSVv1Suite extends CSVSuite { From 54e62a158ead91d832d477a76aace40ef5b54121 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B8rn=20J=C3=B8rgensen?= Date: Thu, 26 Sep 2024 13:37:39 -0700 Subject: [PATCH 150/189] [SPARK-49800][BUILD][K8S] Upgrade `kubernetes-client` to 6.13.4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Upgrade `kubernetes-client` from 6.13.3 to 6.13.4 ### Why are the changes needed? New version that have 5 fixes [Release log 6.13.4](https://github.com/fabric8io/kubernetes-client/releases/tag/v6.13.4) ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48268 from bjornjorgensen/k8sclient6.13.4. Authored-by: Bjørn Jørgensen Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 50 +++++++++++++-------------- pom.xml | 2 +- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 19b8a237d30aa..c9a32757554be 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -159,31 +159,31 @@ jsr305/3.0.0//jsr305-3.0.0.jar jta/1.1//jta-1.1.jar jul-to-slf4j/2.0.16//jul-to-slf4j-2.0.16.jar kryo-shaded/4.0.2//kryo-shaded-4.0.2.jar -kubernetes-client-api/6.13.3//kubernetes-client-api-6.13.3.jar -kubernetes-client/6.13.3//kubernetes-client-6.13.3.jar -kubernetes-httpclient-okhttp/6.13.3//kubernetes-httpclient-okhttp-6.13.3.jar -kubernetes-model-admissionregistration/6.13.3//kubernetes-model-admissionregistration-6.13.3.jar -kubernetes-model-apiextensions/6.13.3//kubernetes-model-apiextensions-6.13.3.jar -kubernetes-model-apps/6.13.3//kubernetes-model-apps-6.13.3.jar -kubernetes-model-autoscaling/6.13.3//kubernetes-model-autoscaling-6.13.3.jar -kubernetes-model-batch/6.13.3//kubernetes-model-batch-6.13.3.jar -kubernetes-model-certificates/6.13.3//kubernetes-model-certificates-6.13.3.jar -kubernetes-model-common/6.13.3//kubernetes-model-common-6.13.3.jar -kubernetes-model-coordination/6.13.3//kubernetes-model-coordination-6.13.3.jar -kubernetes-model-core/6.13.3//kubernetes-model-core-6.13.3.jar -kubernetes-model-discovery/6.13.3//kubernetes-model-discovery-6.13.3.jar -kubernetes-model-events/6.13.3//kubernetes-model-events-6.13.3.jar -kubernetes-model-extensions/6.13.3//kubernetes-model-extensions-6.13.3.jar -kubernetes-model-flowcontrol/6.13.3//kubernetes-model-flowcontrol-6.13.3.jar -kubernetes-model-gatewayapi/6.13.3//kubernetes-model-gatewayapi-6.13.3.jar -kubernetes-model-metrics/6.13.3//kubernetes-model-metrics-6.13.3.jar -kubernetes-model-networking/6.13.3//kubernetes-model-networking-6.13.3.jar -kubernetes-model-node/6.13.3//kubernetes-model-node-6.13.3.jar -kubernetes-model-policy/6.13.3//kubernetes-model-policy-6.13.3.jar -kubernetes-model-rbac/6.13.3//kubernetes-model-rbac-6.13.3.jar -kubernetes-model-resource/6.13.3//kubernetes-model-resource-6.13.3.jar -kubernetes-model-scheduling/6.13.3//kubernetes-model-scheduling-6.13.3.jar -kubernetes-model-storageclass/6.13.3//kubernetes-model-storageclass-6.13.3.jar +kubernetes-client-api/6.13.4//kubernetes-client-api-6.13.4.jar +kubernetes-client/6.13.4//kubernetes-client-6.13.4.jar +kubernetes-httpclient-okhttp/6.13.4//kubernetes-httpclient-okhttp-6.13.4.jar +kubernetes-model-admissionregistration/6.13.4//kubernetes-model-admissionregistration-6.13.4.jar +kubernetes-model-apiextensions/6.13.4//kubernetes-model-apiextensions-6.13.4.jar +kubernetes-model-apps/6.13.4//kubernetes-model-apps-6.13.4.jar +kubernetes-model-autoscaling/6.13.4//kubernetes-model-autoscaling-6.13.4.jar +kubernetes-model-batch/6.13.4//kubernetes-model-batch-6.13.4.jar +kubernetes-model-certificates/6.13.4//kubernetes-model-certificates-6.13.4.jar +kubernetes-model-common/6.13.4//kubernetes-model-common-6.13.4.jar +kubernetes-model-coordination/6.13.4//kubernetes-model-coordination-6.13.4.jar +kubernetes-model-core/6.13.4//kubernetes-model-core-6.13.4.jar +kubernetes-model-discovery/6.13.4//kubernetes-model-discovery-6.13.4.jar +kubernetes-model-events/6.13.4//kubernetes-model-events-6.13.4.jar +kubernetes-model-extensions/6.13.4//kubernetes-model-extensions-6.13.4.jar +kubernetes-model-flowcontrol/6.13.4//kubernetes-model-flowcontrol-6.13.4.jar +kubernetes-model-gatewayapi/6.13.4//kubernetes-model-gatewayapi-6.13.4.jar +kubernetes-model-metrics/6.13.4//kubernetes-model-metrics-6.13.4.jar +kubernetes-model-networking/6.13.4//kubernetes-model-networking-6.13.4.jar +kubernetes-model-node/6.13.4//kubernetes-model-node-6.13.4.jar +kubernetes-model-policy/6.13.4//kubernetes-model-policy-6.13.4.jar +kubernetes-model-rbac/6.13.4//kubernetes-model-rbac-6.13.4.jar +kubernetes-model-resource/6.13.4//kubernetes-model-resource-6.13.4.jar +kubernetes-model-scheduling/6.13.4//kubernetes-model-scheduling-6.13.4.jar +kubernetes-model-storageclass/6.13.4//kubernetes-model-storageclass-6.13.4.jar lapack/3.0.3//lapack-3.0.3.jar leveldbjni-all/1.8//leveldbjni-all-1.8.jar libfb303/0.9.3//libfb303-0.9.3.jar diff --git a/pom.xml b/pom.xml index f3dc92426ac4e..22048b55da27f 100644 --- a/pom.xml +++ b/pom.xml @@ -231,7 +231,7 @@ org.fusesource.leveldbjni - 6.13.3 + 6.13.4 1.17.6 ${java.home} From 339dd5b93316fecd0455b53b2cedee2b5333a184 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 26 Sep 2024 13:39:02 -0700 Subject: [PATCH 151/189] [SPARK-49791][SQL] Make DelegatingCatalogExtension more extendable ### What changes were proposed in this pull request? This PR updates `DelegatingCatalogExtension` so that it's more extendable - `initialize` becomes not final, so that sub-classes can overwrite it - `delegate` becomes `protected`, so that sub-classes can access it In addition, this PR fixes a mistake that `DelegatingCatalogExtension` is just a convenient default implementation, it's actually the `CatalogExtension` interface that indicates this catalog implementation will delegate requests to the Spark session catalog. https://github.com/apache/spark/pull/47724 should use `CatalogExtension` instead. ### Why are the changes needed? Unblock the Iceberg extension. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? existing tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #48257 from cloud-fan/catalog. Lead-authored-by: Wenchen Fan Co-authored-by: Wenchen Fan Signed-off-by: Dongjoon Hyun --- .../sql/connector/catalog/DelegatingCatalogExtension.java | 4 ++-- .../spark/sql/catalyst/analysis/ResolveSessionCatalog.scala | 4 ++-- .../org/apache/spark/sql/internal/DataFrameWriterImpl.scala | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/DelegatingCatalogExtension.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/DelegatingCatalogExtension.java index f6686d2e4d3b6..786821514822e 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/DelegatingCatalogExtension.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/DelegatingCatalogExtension.java @@ -38,7 +38,7 @@ @Evolving public abstract class DelegatingCatalogExtension implements CatalogExtension { - private CatalogPlugin delegate; + protected CatalogPlugin delegate; @Override public final void setDelegateCatalog(CatalogPlugin delegate) { @@ -51,7 +51,7 @@ public String name() { } @Override - public final void initialize(String name, CaseInsensitiveStringMap options) {} + public void initialize(String name, CaseInsensitiveStringMap options) {} @Override public Set capabilities() { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala index 02ad2e79a5645..a9ad7523c8fbc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, toPrettySQL, ResolveDefaultColumns => DefaultCols} import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns._ -import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, CatalogV2Util, DelegatingCatalogExtension, LookupCatalog, SupportsNamespaces, V1Table} +import org.apache.spark.sql.connector.catalog.{CatalogExtension, CatalogManager, CatalogPlugin, CatalogV2Util, LookupCatalog, SupportsNamespaces, V1Table} import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command._ @@ -706,6 +706,6 @@ class ResolveSessionCatalog(val catalogManager: CatalogManager) private def supportsV1Command(catalog: CatalogPlugin): Boolean = { isSessionCatalog(catalog) && ( SQLConf.get.getConf(SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION).isEmpty || - catalog.isInstanceOf[DelegatingCatalogExtension]) + catalog.isInstanceOf[CatalogExtension]) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/DataFrameWriterImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/DataFrameWriterImpl.scala index f0eef9ae1cbb0..8164d33f46fee 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/DataFrameWriterImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/DataFrameWriterImpl.scala @@ -429,7 +429,7 @@ final class DataFrameWriterImpl[T] private[sql](ds: Dataset[T]) extends DataFram val canUseV2 = lookupV2Provider().isDefined || (df.sparkSession.sessionState.conf.getConf( SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION).isDefined && !df.sparkSession.sessionState.catalogManager.catalog(CatalogManager.SESSION_CATALOG_NAME) - .isInstanceOf[DelegatingCatalogExtension]) + .isInstanceOf[CatalogExtension]) session.sessionState.sqlParser.parseMultipartIdentifier(tableName) match { case nameParts @ NonSessionCatalogAndIdentifier(catalog, ident) => From fc9d421a2345987105aa97947c867ac80ba48a05 Mon Sep 17 00:00:00 2001 From: Rui Wang Date: Fri, 27 Sep 2024 08:26:24 +0800 Subject: [PATCH 152/189] [SPARK-49211][SQL][FOLLOW-UP] Support catalog in QualifiedTableName ### What changes were proposed in this pull request? Support catalog in QualifiedTableName and remove `FullQualifiedTableName`. ### Why are the changes needed? Consolidate and remove duplicate code. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing UT ### Was this patch authored or co-authored using generative AI tooling? No Closes #48255 from amaliujia/qualifedtablename. Authored-by: Rui Wang Signed-off-by: Wenchen Fan --- .../sql/catalyst/catalog/SessionCatalog.scala | 18 ++++++++-------- .../spark/sql/catalyst/identifiers.scala | 21 +++++++++++++++---- .../catalog/SessionCatalogSuite.scala | 8 +++---- .../datasources/DataSourceStrategy.scala | 4 ++-- .../datasources/v2/V2SessionCatalog.scala | 4 ++-- .../sql/StatisticsCollectionTestBase.scala | 4 ++-- .../sql/connector/DataSourceV2SQLSuite.scala | 10 ++++----- .../sql/execution/command/DDLSuite.scala | 6 +++--- .../command/v1/TruncateTableSuite.scala | 4 ++-- .../spark/sql/hive/HiveMetastoreCatalog.scala | 12 +++++------ 10 files changed, 52 insertions(+), 39 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala index d3a6cb6ae2845..a0f7af10fefaf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala @@ -197,7 +197,7 @@ class SessionCatalog( } } - private val tableRelationCache: Cache[FullQualifiedTableName, LogicalPlan] = { + private val tableRelationCache: Cache[QualifiedTableName, LogicalPlan] = { var builder = CacheBuilder.newBuilder() .maximumSize(cacheSize) @@ -205,33 +205,33 @@ class SessionCatalog( builder = builder.expireAfterWrite(cacheTTL, TimeUnit.SECONDS) } - builder.build[FullQualifiedTableName, LogicalPlan]() + builder.build[QualifiedTableName, LogicalPlan]() } /** This method provides a way to get a cached plan. */ - def getCachedPlan(t: FullQualifiedTableName, c: Callable[LogicalPlan]): LogicalPlan = { + def getCachedPlan(t: QualifiedTableName, c: Callable[LogicalPlan]): LogicalPlan = { tableRelationCache.get(t, c) } /** This method provides a way to get a cached plan if the key exists. */ - def getCachedTable(key: FullQualifiedTableName): LogicalPlan = { + def getCachedTable(key: QualifiedTableName): LogicalPlan = { tableRelationCache.getIfPresent(key) } /** This method provides a way to cache a plan. */ - def cacheTable(t: FullQualifiedTableName, l: LogicalPlan): Unit = { + def cacheTable(t: QualifiedTableName, l: LogicalPlan): Unit = { tableRelationCache.put(t, l) } /** This method provides a way to invalidate a cached plan. */ - def invalidateCachedTable(key: FullQualifiedTableName): Unit = { + def invalidateCachedTable(key: QualifiedTableName): Unit = { tableRelationCache.invalidate(key) } /** This method discards any cached table relation plans for the given table identifier. */ def invalidateCachedTable(name: TableIdentifier): Unit = { val qualified = qualifyIdentifier(name) - invalidateCachedTable(FullQualifiedTableName( + invalidateCachedTable(QualifiedTableName( qualified.catalog.get, qualified.database.get, qualified.table)) } @@ -301,7 +301,7 @@ class SessionCatalog( } if (cascade && databaseExists(dbName)) { listTables(dbName).foreach { t => - invalidateCachedTable(FullQualifiedTableName(SESSION_CATALOG_NAME, dbName, t.table)) + invalidateCachedTable(QualifiedTableName(SESSION_CATALOG_NAME, dbName, t.table)) } } externalCatalog.dropDatabase(dbName, ignoreIfNotExists, cascade) @@ -1183,7 +1183,7 @@ class SessionCatalog( def refreshTable(name: TableIdentifier): Unit = synchronized { getLocalOrGlobalTempView(name).map(_.refresh()).getOrElse { val qualifiedIdent = qualifyIdentifier(name) - val qualifiedTableName = FullQualifiedTableName( + val qualifiedTableName = QualifiedTableName( qualifiedIdent.catalog.get, qualifiedIdent.database.get, qualifiedIdent.table) tableRelationCache.invalidate(qualifiedTableName) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/identifiers.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/identifiers.scala index cc881539002b6..ceced9313940a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/identifiers.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/identifiers.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.catalyst +import org.apache.spark.sql.connector.catalog.CatalogManager + /** * An identifier that optionally specifies a database. * @@ -107,14 +109,25 @@ case class TableIdentifier(table: String, database: Option[String], catalog: Opt } /** A fully qualified identifier for a table (i.e., database.tableName) */ -case class QualifiedTableName(database: String, name: String) { - override def toString: String = s"$database.$name" -} +case class QualifiedTableName(catalog: String, database: String, name: String) { + /** Two argument ctor for backward compatibility. */ + def this(database: String, name: String) = this( + catalog = CatalogManager.SESSION_CATALOG_NAME, + database = database, + name = name) -case class FullQualifiedTableName(catalog: String, database: String, name: String) { override def toString: String = s"$catalog.$database.$name" } +object QualifiedTableName { + def apply(catalog: String, database: String, name: String): QualifiedTableName = { + new QualifiedTableName(catalog, database, name) + } + + def apply(database: String, name: String): QualifiedTableName = + new QualifiedTableName(database = database, name = name) +} + object TableIdentifier { def apply(tableName: String): TableIdentifier = new TableIdentifier(tableName) def apply(table: String, database: Option[String]): TableIdentifier = diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalogSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalogSuite.scala index fbe63f71ae029..cfbc507fb5c74 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalogSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalogSuite.scala @@ -22,7 +22,7 @@ import scala.concurrent.duration._ import org.scalatest.concurrent.Eventually import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.{AliasIdentifier, FullQualifiedTableName, FunctionIdentifier, TableIdentifier} +import org.apache.spark.sql.catalyst.{AliasIdentifier, FunctionIdentifier, QualifiedTableName, TableIdentifier} import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser @@ -1883,7 +1883,7 @@ abstract class SessionCatalogSuite extends AnalysisTest with Eventually { conf.setConf(StaticSQLConf.METADATA_CACHE_TTL_SECONDS, 1L) withConfAndEmptyCatalog(conf) { catalog => - val table = FullQualifiedTableName( + val table = QualifiedTableName( CatalogManager.SESSION_CATALOG_NAME, catalog.getCurrentDatabase, "test") // First, make sure the test table is not cached. @@ -1903,14 +1903,14 @@ abstract class SessionCatalogSuite extends AnalysisTest with Eventually { test("SPARK-34197: refreshTable should not invalidate the relation cache for temporary views") { withBasicCatalog { catalog => createTempView(catalog, "tbl1", Range(1, 10, 1, 10), false) - val qualifiedName1 = FullQualifiedTableName(SESSION_CATALOG_NAME, "default", "tbl1") + val qualifiedName1 = QualifiedTableName(SESSION_CATALOG_NAME, "default", "tbl1") catalog.cacheTable(qualifiedName1, Range(1, 10, 1, 10)) catalog.refreshTable(TableIdentifier("tbl1")) assert(catalog.getCachedTable(qualifiedName1) != null) createGlobalTempView(catalog, "tbl2", Range(2, 10, 1, 10), false) val qualifiedName2 = - FullQualifiedTableName(SESSION_CATALOG_NAME, catalog.globalTempDatabase, "tbl2") + QualifiedTableName(SESSION_CATALOG_NAME, catalog.globalTempDatabase, "tbl2") catalog.cacheTable(qualifiedName2, Range(2, 10, 1, 10)) catalog.refreshTable(TableIdentifier("tbl2", Some(catalog.globalTempDatabase))) assert(catalog.getCachedTable(qualifiedName2) != null) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala index 2be4b236872f0..a2707da2d1023 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala @@ -28,7 +28,7 @@ import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys.PREDICATES import org.apache.spark.rdd.RDD import org.apache.spark.sql._ -import org.apache.spark.sql.catalyst.{expressions, CatalystTypeConverters, FullQualifiedTableName, InternalRow, SQLConfHelper} +import org.apache.spark.sql.catalyst.{expressions, CatalystTypeConverters, InternalRow, QualifiedTableName, SQLConfHelper} import org.apache.spark.sql.catalyst.CatalystTypeConverters.convertToScala import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.catalog._ @@ -249,7 +249,7 @@ class FindDataSourceTable(sparkSession: SparkSession) extends Rule[LogicalPlan] private def readDataSourceTable( table: CatalogTable, extraOptions: CaseInsensitiveStringMap): LogicalPlan = { val qualifiedTableName = - FullQualifiedTableName(table.identifier.catalog.get, table.database, table.identifier.table) + QualifiedTableName(table.identifier.catalog.get, table.database, table.identifier.table) val catalog = sparkSession.sessionState.catalog val dsOptions = DataSourceUtils.generateDatasourceOptions(extraOptions, table) catalog.getCachedPlan(qualifiedTableName, () => { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2SessionCatalog.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2SessionCatalog.scala index bd1df87d15c3c..22c13fd98ced1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2SessionCatalog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2SessionCatalog.scala @@ -24,7 +24,7 @@ import scala.collection.mutable import scala.jdk.CollectionConverters._ import org.apache.spark.SparkUnsupportedOperationException -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, FunctionIdentifier, SQLConfHelper, TableIdentifier} +import org.apache.spark.sql.catalyst.{FunctionIdentifier, QualifiedTableName, SQLConfHelper, TableIdentifier} import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException, TableAlreadyExistsException} import org.apache.spark.sql.catalyst.catalog.{CatalogDatabase, CatalogStorageFormat, CatalogTable, CatalogTableType, CatalogUtils, ClusterBySpec, SessionCatalog} import org.apache.spark.sql.catalyst.util.TypeUtils._ @@ -93,7 +93,7 @@ class V2SessionCatalog(catalog: SessionCatalog) // table here. To avoid breaking it we do not resolve the table provider and still return // `V1Table` if the custom session catalog is present. if (table.provider.isDefined && !hasCustomSessionCatalog) { - val qualifiedTableName = FullQualifiedTableName( + val qualifiedTableName = QualifiedTableName( table.identifier.catalog.get, table.database, table.identifier.table) // Check if the table is in the v1 table cache to skip the v2 table lookup. if (catalog.getCachedTable(qualifiedTableName) != null) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala index 7fa29dd38fd96..74329ac0e0d23 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionTestBase.scala @@ -25,7 +25,7 @@ import java.time.LocalDateTime import scala.collection.mutable import scala.util.Random -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, TableIdentifier} +import org.apache.spark.sql.catalyst.{QualifiedTableName, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics, CatalogTable, HiveTableRelation} import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec import org.apache.spark.sql.catalyst.expressions.AttributeMap @@ -270,7 +270,7 @@ abstract class StatisticsCollectionTestBase extends QueryTest with SQLTestUtils def getTableFromCatalogCache(tableName: String): LogicalPlan = { val catalog = spark.sessionState.catalog - val qualifiedTableName = FullQualifiedTableName( + val qualifiedTableName = QualifiedTableName( CatalogManager.SESSION_CATALOG_NAME, catalog.getCurrentDatabase, tableName) catalog.getCachedTable(qualifiedTableName) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index 7aaec6d500ba0..dac066bbef838 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.{SparkException, SparkRuntimeException, SparkUnsupportedOperationException} import org.apache.spark.sql._ -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, InternalRow, TableIdentifier} +import org.apache.spark.sql.catalyst.{InternalRow, QualifiedTableName, TableIdentifier} import org.apache.spark.sql.catalyst.CurrentUserContext.CURRENT_USER import org.apache.spark.sql.catalyst.analysis.{CannotReplaceMissingTableException, NoSuchNamespaceException, TableAlreadyExistsException} import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType, CatalogUtils} @@ -3713,7 +3713,7 @@ class DataSourceV2SQLSuiteV1Filter // Reset CatalogManager to clear the materialized `spark_catalog` instance, so that we can // configure a new implementation. - val table1 = FullQualifiedTableName(SESSION_CATALOG_NAME, "default", "t") + val table1 = QualifiedTableName(SESSION_CATALOG_NAME, "default", "t") spark.sessionState.catalogManager.reset() withSQLConf( V2_SESSION_CATALOG_IMPLEMENTATION.key -> @@ -3722,7 +3722,7 @@ class DataSourceV2SQLSuiteV1Filter checkParquet(table1.toString, path.getAbsolutePath) } } - val table2 = FullQualifiedTableName("testcat3", "default", "t") + val table2 = QualifiedTableName("testcat3", "default", "t") withSQLConf( "spark.sql.catalog.testcat3" -> classOf[V2CatalogSupportBuiltinDataSource].getName) { withTempPath { path => @@ -3741,7 +3741,7 @@ class DataSourceV2SQLSuiteV1Filter // Reset CatalogManager to clear the materialized `spark_catalog` instance, so that we can // configure a new implementation. spark.sessionState.catalogManager.reset() - val table1 = FullQualifiedTableName(SESSION_CATALOG_NAME, "default", "t") + val table1 = QualifiedTableName(SESSION_CATALOG_NAME, "default", "t") withSQLConf( V2_SESSION_CATALOG_IMPLEMENTATION.key -> classOf[V2CatalogSupportBuiltinDataSource].getName) { @@ -3750,7 +3750,7 @@ class DataSourceV2SQLSuiteV1Filter } } - val table2 = FullQualifiedTableName("testcat3", "default", "t") + val table2 = QualifiedTableName("testcat3", "default", "t") withSQLConf( "spark.sql.catalog.testcat3" -> classOf[V2CatalogSupportBuiltinDataSource].getName) { withTempPath { path => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala index 8307326f17fcf..e07f6406901e0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLSuite.scala @@ -27,7 +27,7 @@ import org.apache.hadoop.fs.permission.{AclEntry, AclStatus} import org.apache.spark.{SparkClassNotFoundException, SparkException, SparkFiles, SparkRuntimeException} import org.apache.spark.internal.config import org.apache.spark.sql.{AnalysisException, QueryTest, Row, SaveMode} -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, FunctionIdentifier, TableIdentifier} +import org.apache.spark.sql.catalyst.{FunctionIdentifier, QualifiedTableName, TableIdentifier} import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException import org.apache.spark.sql.catalyst.catalog._ import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec @@ -219,7 +219,7 @@ class InMemoryCatalogedDDLSuite extends DDLSuite with SharedSparkSession { test("SPARK-25403 refresh the table after inserting data") { withTable("t") { val catalog = spark.sessionState.catalog - val table = FullQualifiedTableName( + val table = QualifiedTableName( CatalogManager.SESSION_CATALOG_NAME, catalog.getCurrentDatabase, "t") sql("CREATE TABLE t (a INT) USING parquet") sql("INSERT INTO TABLE t VALUES (1)") @@ -233,7 +233,7 @@ class InMemoryCatalogedDDLSuite extends DDLSuite with SharedSparkSession { withTable("t") { withTempDir { dir => val catalog = spark.sessionState.catalog - val table = FullQualifiedTableName( + val table = QualifiedTableName( CatalogManager.SESSION_CATALOG_NAME, catalog.getCurrentDatabase, "t") val p1 = s"${dir.getCanonicalPath}/p1" val p2 = s"${dir.getCanonicalPath}/p2" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/TruncateTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/TruncateTableSuite.scala index 348b216aeb044..40ae35bbe8aa3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/TruncateTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/TruncateTableSuite.scala @@ -23,7 +23,7 @@ import org.apache.hadoop.fs.Path import org.apache.hadoop.fs.permission.{AclEntry, AclEntryScope, AclEntryType, FsAction, FsPermission} import org.apache.spark.sql.{AnalysisException, Row} -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, TableIdentifier} +import org.apache.spark.sql.catalyst.{QualifiedTableName, TableIdentifier} import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.execution.command import org.apache.spark.sql.execution.command.FakeLocalFsFileSystem @@ -148,7 +148,7 @@ trait TruncateTableSuiteBase extends command.TruncateTableSuiteBase { val catalog = spark.sessionState.catalog val qualifiedTableName = - FullQualifiedTableName(CatalogManager.SESSION_CATALOG_NAME, "ns", "tbl") + QualifiedTableName(CatalogManager.SESSION_CATALOG_NAME, "ns", "tbl") val cachedPlan = catalog.getCachedTable(qualifiedTableName) assert(cachedPlan.stats.sizeInBytes == 0) } diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala index 7873c36222da0..1f87db31ffa52 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala @@ -28,7 +28,7 @@ import org.apache.spark.SparkException import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys._ import org.apache.spark.sql.{AnalysisException, SparkSession} -import org.apache.spark.sql.catalyst.{FullQualifiedTableName, TableIdentifier} +import org.apache.spark.sql.catalyst.{QualifiedTableName, TableIdentifier} import org.apache.spark.sql.catalyst.catalog._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.types.DataTypeUtils @@ -56,7 +56,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log private val tableCreationLocks = Striped.lazyWeakLock(100) /** Acquires a lock on the table cache for the duration of `f`. */ - private def withTableCreationLock[A](tableName: FullQualifiedTableName, f: => A): A = { + private def withTableCreationLock[A](tableName: QualifiedTableName, f: => A): A = { val lock = tableCreationLocks.get(tableName) lock.lock() try f finally { @@ -66,7 +66,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log // For testing only private[hive] def getCachedDataSourceTable(table: TableIdentifier): LogicalPlan = { - val key = FullQualifiedTableName( + val key = QualifiedTableName( // scalastyle:off caselocale table.catalog.getOrElse(CatalogManager.SESSION_CATALOG_NAME).toLowerCase, table.database.getOrElse(sessionState.catalog.getCurrentDatabase).toLowerCase, @@ -76,7 +76,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log } private def getCached( - tableIdentifier: FullQualifiedTableName, + tableIdentifier: QualifiedTableName, pathsInMetastore: Seq[Path], schemaInMetastore: StructType, expectedFileFormat: Class[_ <: FileFormat], @@ -120,7 +120,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log } private def logWarningUnexpectedFileFormat( - tableIdentifier: FullQualifiedTableName, + tableIdentifier: QualifiedTableName, expectedFileFormat: Class[_ <: FileFormat], actualFileFormat: String): Unit = { logWarning(log"Table ${MDC(TABLE_NAME, tableIdentifier)} should be stored as " + @@ -201,7 +201,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log fileType: String, isWrite: Boolean): LogicalRelation = { val metastoreSchema = relation.tableMeta.schema - val tableIdentifier = FullQualifiedTableName(relation.tableMeta.identifier.catalog.get, + val tableIdentifier = QualifiedTableName(relation.tableMeta.identifier.catalog.get, relation.tableMeta.database, relation.tableMeta.identifier.table) val lazyPruningEnabled = sparkSession.sessionState.conf.manageFilesourcePartitions From 09b7aa67ce64d7d4ecc803215eaf85464df181c5 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Thu, 26 Sep 2024 17:32:29 -0700 Subject: [PATCH 153/189] [SPARK-49803][SQL][TESTS] Increase `spark.test.docker.connectionTimeout` to 10min ### What changes were proposed in this pull request? This PR aims to increase `spark.test.docker.connectionTimeout` to 10min. ### Why are the changes needed? Recently, various DB images fails at `connection` stage on multiple branches. **MASTER** branch https://github.com/apache/spark/actions/runs/11045311764/job/30682732260 ``` [info] OracleIntegrationSuite: [info] org.apache.spark.sql.jdbc.OracleIntegrationSuite *** ABORTED *** (5 minutes, 17 seconds) [info] The code passed to eventually never returned normally. Attempted 298 times over 5.0045005511500005 minutes. Last failure message: ORA-12541: Cannot connect. No listener at host 10.1.0.41 port 41079. (CONNECTION_ID=n9ZWIh+nQn+G9fkwKyoBQA==) ``` **branch-3.5** branch https://github.com/apache/spark/actions/runs/10939696926/job/30370552237 ``` [info] MsSqlServerNamespaceSuite: [info] org.apache.spark.sql.jdbc.v2.MsSqlServerNamespaceSuite *** ABORTED *** (5 minutes, 42 seconds) [info] The code passed to eventually never returned normally. Attempted 11 times over 5.487631282400001 minutes. Last failure message: The TCP/IP connection to the host 10.1.0.56, port 35345 has failed. Error: "Connection refused (Connection refused). Verify the connection properties. Make sure that an instance of SQL Server is running on the host and accepting TCP/IP connections at the port. Make sure that TCP connections to the port are not blocked by a firewall.".. (DockerJDBCIntegrationSuite.scala:166) ``` **branch-3.4** branch https://github.com/apache/spark/actions/runs/10937842509/job/30364658576 ``` [info] MsSqlServerNamespaceSuite: [info] org.apache.spark.sql.jdbc.v2.MsSqlServerNamespaceSuite *** ABORTED *** (5 minutes, 42 seconds) [info] The code passed to eventually never returned normally. Attempted 11 times over 5.487555645633333 minutes. Last failure message: The TCP/IP connection to the host 10.1.0.153, port 46153 has failed. Error: "Connection refused (Connection refused). Verify the connection properties. Make sure that an instance of SQL Server is running on the host and accepting TCP/IP connections at the port. Make sure that TCP connections to the port are not blocked by a firewall.".. (DockerJDBCIntegrationSuite.scala:166) ``` ### Does this PR introduce _any_ user-facing change? No, this is a test-only change. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48272 from dongjoon-hyun/SPARK-49803. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala index 8d17e0b4e36e6..1df01bd3bfb62 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala @@ -115,7 +115,7 @@ abstract class DockerJDBCIntegrationSuite protected val startContainerTimeout: Long = timeStringAsSeconds(sys.props.getOrElse("spark.test.docker.startContainerTimeout", "5min")) protected val connectionTimeout: PatienceConfiguration.Timeout = { - val timeoutStr = sys.props.getOrElse("spark.test.docker.connectionTimeout", "5min") + val timeoutStr = sys.props.getOrElse("spark.test.docker.connectionTimeout", "10min") timeout(timeStringAsSeconds(timeoutStr).seconds) } From 488c3f604490c8632dde67a00118d49ccfcbf578 Mon Sep 17 00:00:00 2001 From: Xinrong Meng Date: Fri, 27 Sep 2024 08:35:10 +0800 Subject: [PATCH 154/189] [SPARK-49776][PYTHON][CONNECT] Support pie plots ### What changes were proposed in this pull request? Support area plots with plotly backend on both Spark Connect and Spark classic. ### Why are the changes needed? While Pandas on Spark supports plotting, PySpark currently lacks this feature. The proposed API will enable users to generate visualizations. This will provide users with an intuitive, interactive way to explore and understand large datasets directly from PySpark DataFrames, streamlining the data analysis workflow in distributed environments. See more at [PySpark Plotting API Specification](https://docs.google.com/document/d/1IjOEzC8zcetG86WDvqkereQPj_NGLNW7Bdu910g30Dg/edit?usp=sharing) in progress. Part of https://issues.apache.org/jira/browse/SPARK-49530. ### Does this PR introduce _any_ user-facing change? Yes. Area plots are supported as shown below. ```py >>> from datetime import datetime >>> data = [ ... (3, 5, 20, datetime(2018, 1, 31)), ... (2, 5, 42, datetime(2018, 2, 28)), ... (3, 6, 28, datetime(2018, 3, 31)), ... (9, 12, 62, datetime(2018, 4, 30))] >>> columns = ["sales", "signups", "visits", "date"] >>> df = spark.createDataFrame(data, columns) >>> fig = df.plot(kind="pie", x="date", y="sales") # df.plot(kind="pie", x="date", y="sales") >>> fig.show() ``` ![newplot (8)](https://github.com/user-attachments/assets/c4078bb7-4d84-4607-bcd7-bdd6fbbf8e28) ### How was this patch tested? Unit tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48256 from xinrong-meng/plot_pie. Authored-by: Xinrong Meng Signed-off-by: Xinrong Meng --- python/pyspark/errors/error-conditions.json | 5 +++ python/pyspark/sql/plot/core.py | 41 ++++++++++++++++++- python/pyspark/sql/plot/plotly.py | 15 +++++++ .../sql/tests/plot/test_frame_plot_plotly.py | 25 +++++++++++ 4 files changed, 85 insertions(+), 1 deletion(-) diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 115ad658e32f5..ed62ea117d369 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -812,6 +812,11 @@ "Pipe function `` exited with error code ." ] }, + "PLOT_NOT_NUMERIC_COLUMN": { + "message": [ + "Argument must be a numerical column for plotting, got ." + ] + }, "PYTHON_HASH_SEED_NOT_SET": { "message": [ "Randomness of hash of string should be disabled via PYTHONHASHSEED." diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index 9f83d00696524..f9667ee2c0d69 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -17,7 +17,8 @@ from typing import Any, TYPE_CHECKING, Optional, Union from types import ModuleType -from pyspark.errors import PySparkRuntimeError, PySparkValueError +from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError +from pyspark.sql.types import NumericType from pyspark.sql.utils import require_minimum_plotly_version @@ -97,6 +98,7 @@ class PySparkPlotAccessor: "bar": PySparkTopNPlotBase().get_top_n, "barh": PySparkTopNPlotBase().get_top_n, "line": PySparkSampledPlotBase().get_sampled, + "pie": PySparkTopNPlotBase().get_top_n, "scatter": PySparkSampledPlotBase().get_sampled, } _backends = {} # type: ignore[var-annotated] @@ -299,3 +301,40 @@ def area(self, x: str, y: str, **kwargs: Any) -> "Figure": >>> df.plot.area(x='date', y=['sales', 'signups', 'visits']) # doctest: +SKIP """ return self(kind="area", x=x, y=y, **kwargs) + + def pie(self, x: str, y: str, **kwargs: Any) -> "Figure": + """ + Generate a pie plot. + + A pie plot is a proportional representation of the numerical data in a + column. + + Parameters + ---------- + x : str + Name of column to be used as the category labels for the pie plot. + y : str + Name of the column to plot. + **kwargs + Additional keyword arguments. + + Returns + ------- + :class:`plotly.graph_objs.Figure` + + Examples + -------- + """ + schema = self.data.schema + + # Check if 'y' is a numerical column + y_field = schema[y] if y in schema.names else None + if y_field is None or not isinstance(y_field.dataType, NumericType): + raise PySparkTypeError( + errorClass="PLOT_NOT_NUMERIC_COLUMN", + messageParameters={ + "arg_name": "y", + "arg_type": str(y_field.dataType) if y_field else "None", + }, + ) + return self(kind="pie", x=x, y=y, **kwargs) diff --git a/python/pyspark/sql/plot/plotly.py b/python/pyspark/sql/plot/plotly.py index 5efc19476057f..91f5363464717 100644 --- a/python/pyspark/sql/plot/plotly.py +++ b/python/pyspark/sql/plot/plotly.py @@ -27,4 +27,19 @@ def plot_pyspark(data: "DataFrame", kind: str, **kwargs: Any) -> "Figure": import plotly + if kind == "pie": + return plot_pie(data, **kwargs) + return plotly.plot(PySparkPlotAccessor.plot_data_map[kind](data), kind, **kwargs) + + +def plot_pie(data: "DataFrame", **kwargs: Any) -> "Figure": + # TODO(SPARK-49530): Support pie subplots with plotly backend + from plotly import express + + pdf = PySparkPlotAccessor.plot_data_map["pie"](data) + x = kwargs.pop("x", None) + y = kwargs.pop("y", None) + fig = express.pie(pdf, values=y, names=x, **kwargs) + + return fig diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py index 6176525b49550..70a1b336f734a 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -19,6 +19,7 @@ from datetime import datetime import pyspark.sql.plot # noqa: F401 +from pyspark.errors import PySparkTypeError from pyspark.testing.sqlutils import ReusedSQLTestCase, have_plotly, plotly_requirement_message @@ -64,6 +65,11 @@ def _check_fig_data(self, kind, fig_data, expected_x, expected_y, expected_name= self.assertEqual(fig_data["type"], "scatter") self.assertEqual(fig_data["orientation"], "v") self.assertEqual(fig_data["mode"], "lines") + elif kind == "pie": + self.assertEqual(fig_data["type"], "pie") + self.assertEqual(list(fig_data["labels"]), expected_x) + self.assertEqual(list(fig_data["values"]), expected_y) + return self.assertEqual(fig_data["xaxis"], "x") self.assertEqual(list(fig_data["x"]), expected_x) @@ -133,6 +139,25 @@ def test_area_plot(self): self._check_fig_data("area", fig["data"][1], expected_x, [5, 5, 6, 12], "signups") self._check_fig_data("area", fig["data"][2], expected_x, [20, 42, 28, 62], "visits") + def test_pie_plot(self): + fig = self.sdf3.plot(kind="pie", x="date", y="sales") + expected_x = [ + datetime(2018, 1, 31, 0, 0), + datetime(2018, 2, 28, 0, 0), + datetime(2018, 3, 31, 0, 0), + datetime(2018, 4, 30, 0, 0), + ] + self._check_fig_data("pie", fig["data"][0], expected_x, [3, 2, 3, 9]) + + # y is not a numerical column + with self.assertRaises(PySparkTypeError) as pe: + self.sdf.plot.pie(x="int_val", y="category") + self.check_error( + exception=pe.exception, + errorClass="PLOT_NOT_NUMERIC_COLUMN", + messageParameters={"arg_name": "y", "arg_type": "StringType()"}, + ) + class DataFramePlotPlotlyTests(DataFramePlotPlotlyTestsMixin, ReusedSQLTestCase): pass From 27d4a77f2a8ccdbc4d7c3afd6743ec845dc1294b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B8rn=20J=C3=B8rgensen?= Date: Fri, 27 Sep 2024 11:33:03 +0900 Subject: [PATCH 155/189] [SPARK-49801][PYTHON][PS][BUILD] Update `pandas` to 2.2.3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Update pandas from 2.2.2 to 2.2.3 ### Why are the changes needed? [Release notes](https://pandas.pydata.org/pandas-docs/version/2.2.3/whatsnew/v2.2.3.html) ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48269 from bjornjorgensen/pandas2.2.3. Authored-by: Bjørn Jørgensen Signed-off-by: Hyukjin Kwon --- dev/infra/Dockerfile | 4 ++-- python/pyspark/pandas/supported_api_gen.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dev/infra/Dockerfile b/dev/infra/Dockerfile index 5939e429b2f35..a40e43bb659f8 100644 --- a/dev/infra/Dockerfile +++ b/dev/infra/Dockerfile @@ -91,10 +91,10 @@ RUN mkdir -p /usr/local/pypy/pypy3.9 && \ ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3.9 && \ ln -sf /usr/local/pypy/pypy3.9/bin/pypy /usr/local/bin/pypy3 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | pypy3 -RUN pypy3 -m pip install 'numpy==1.26.4' 'six==1.16.0' 'pandas==2.2.2' scipy coverage matplotlib lxml +RUN pypy3 -m pip install 'numpy==1.26.4' 'six==1.16.0' 'pandas==2.2.3' scipy coverage matplotlib lxml -ARG BASIC_PIP_PKGS="numpy==1.26.4 pyarrow>=15.0.0 six==1.16.0 pandas==2.2.2 scipy plotly>=4.8 mlflow>=2.8.1 coverage matplotlib openpyxl memory-profiler>=0.61.0 scikit-learn>=1.3.2" +ARG BASIC_PIP_PKGS="numpy==1.26.4 pyarrow>=15.0.0 six==1.16.0 pandas==2.2.3 scipy plotly>=4.8 mlflow>=2.8.1 coverage matplotlib openpyxl memory-profiler>=0.61.0 scikit-learn>=1.3.2" # Python deps for Spark Connect ARG CONNECT_PIP_PKGS="grpcio==1.62.0 grpcio-status==1.62.0 protobuf==4.25.1 googleapis-common-protos==1.56.4 graphviz==0.20.3" diff --git a/python/pyspark/pandas/supported_api_gen.py b/python/pyspark/pandas/supported_api_gen.py index bbf0b3cbc3d67..f2a73cb1c1adf 100644 --- a/python/pyspark/pandas/supported_api_gen.py +++ b/python/pyspark/pandas/supported_api_gen.py @@ -38,7 +38,7 @@ MAX_MISSING_PARAMS_SIZE = 5 COMMON_PARAMETER_SET = {"kwargs", "args", "cls"} MODULE_GROUP_MATCH = [(pd, ps), (pdw, psw), (pdg, psg)] -PANDAS_LATEST_VERSION = "2.2.2" +PANDAS_LATEST_VERSION = "2.2.3" RST_HEADER = """ ===================== From 5d701f2d5add05b7af3889d6b87a192c11872298 Mon Sep 17 00:00:00 2001 From: "oleksii.diagiliev" Date: Thu, 26 Sep 2024 21:59:12 -0700 Subject: [PATCH 156/189] [SPARK-49804][K8S] Fix to use the exit code of executor container always ### What changes were proposed in this pull request? When deploying Spark pods on Kubernetes with sidecars, the reported executor's exit code may be incorrect. For example, the reported executor's exit code is 0(success), but the actual is 52 (OOM). ``` 2024-09-25 02:35:29,383 ERROR TaskSchedulerImpl.logExecutorLoss - Lost executor 1 on XXXXX: The executor with id 1 exited with exit code 0(success). The API gave the following container statuses: container name: fluentd container image: docker-images-release.XXXXX.com/XXXXX/fluentd:XXXXX container state: terminated container started at: 2024-09-25T02:32:17Z container finished at: 2024-09-25T02:34:52Z exit code: 0 termination reason: Completed container name: istio-proxy container image: docker-images-release.XXXXX.com/XXXXX-istio/proxyv2:XXXXX container state: running container started at: 2024-09-25T02:32:16Z container name: spark-kubernetes-executor container image: docker-dev-artifactory.XXXXX.com/XXXXX/spark-XXXXX:XXXXX container state: terminated container started at: 2024-09-25T02:32:17Z container finished at: 2024-09-25T02:35:28Z exit code: 52 termination reason: Error ``` The `ExecutorPodsLifecycleManager.findExitCode()` looks for any terminated container and may choose the sidecar instead of the main executor container. I'm changing it to look for the executor container always. Note, it may happen that the pod fails because of the failure of the sidecar container while executor's container is still running, with my changes the reported exit code will be -1 (`UNKNOWN_EXIT_CODE`). ### Why are the changes needed? To correctly report executor failure reason on UI, in the logs and for the event listeners `SparkListener.onExecutorRemoved()` ### Does this PR introduce _any_ user-facing change? Yes, the executor's exit code is taken from the main container instead of the sidecar. ### How was this patch tested? Added unit test and tested manually on the Kubernetes cluster by simulating different types of executor failure (JVM OOM and container eviction due to disk pressure on the node). ### Was this patch authored or co-authored using generative AI tooling? No Closes #48275 from fe2s/SPARK-49804-fix-exit-code. Lead-authored-by: oleksii.diagiliev Co-authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../k8s/ExecutorPodsLifecycleManager.scala | 6 ++- .../k8s/ExecutorLifecycleTestUtils.scala | 37 ++++++++++++++++++- .../ExecutorPodsLifecycleManagerSuite.scala | 14 ++++++- 3 files changed, 53 insertions(+), 4 deletions(-) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala index 0d79efa06e497..992be9099639e 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala @@ -62,6 +62,9 @@ private[spark] class ExecutorPodsLifecycleManager( private val namespace = conf.get(KUBERNETES_NAMESPACE) + private val sparkContainerName = conf.get(KUBERNETES_EXECUTOR_PODTEMPLATE_CONTAINER_NAME) + .getOrElse(DEFAULT_EXECUTOR_CONTAINER_NAME) + def start(schedulerBackend: KubernetesClusterSchedulerBackend): Unit = { val eventProcessingInterval = conf.get(KUBERNETES_EXECUTOR_EVENT_PROCESSING_INTERVAL) snapshotsStore.addSubscriber(eventProcessingInterval) { @@ -246,7 +249,8 @@ private[spark] class ExecutorPodsLifecycleManager( private def findExitCode(podState: FinalPodState): Int = { podState.pod.getStatus.getContainerStatuses.asScala.find { containerStatus => - containerStatus.getState.getTerminated != null + containerStatus.getName == sparkContainerName && + containerStatus.getState.getTerminated != null }.map { terminatedContainer => terminatedContainer.getState.getTerminated.getExitCode.toInt }.getOrElse(UNKNOWN_EXIT_CODE) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorLifecycleTestUtils.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorLifecycleTestUtils.scala index 299979071b5d7..fc75414e4a7e0 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorLifecycleTestUtils.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorLifecycleTestUtils.scala @@ -29,6 +29,7 @@ import org.apache.spark.resource.ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID object ExecutorLifecycleTestUtils { val TEST_SPARK_APP_ID = "spark-app-id" + val TEST_SPARK_EXECUTOR_CONTAINER_NAME = "spark-executor" def failedExecutorWithoutDeletion( executorId: Long, rpId: Int = DEFAULT_RESOURCE_PROFILE_ID): Pod = { @@ -37,7 +38,7 @@ object ExecutorLifecycleTestUtils { .withPhase("failed") .withStartTime(Instant.now.toString) .addNewContainerStatus() - .withName("spark-executor") + .withName(TEST_SPARK_EXECUTOR_CONTAINER_NAME) .withImage("k8s-spark") .withNewState() .withNewTerminated() @@ -49,6 +50,38 @@ object ExecutorLifecycleTestUtils { .addNewContainerStatus() .withName("spark-executor-sidecar") .withImage("k8s-spark-sidecar") + .withNewState() + .withNewTerminated() + .withMessage("Failed") + .withExitCode(2) + .endTerminated() + .endState() + .endContainerStatus() + .withMessage("Executor failed.") + .withReason("Executor failed because of a thrown error.") + .endStatus() + .build() + } + + def failedExecutorWithSidecarStatusListedFirst( + executorId: Long, rpId: Int = DEFAULT_RESOURCE_PROFILE_ID): Pod = { + new PodBuilder(podWithAttachedContainerForId(executorId, rpId)) + .editOrNewStatus() + .withPhase("failed") + .withStartTime(Instant.now.toString) + .addNewContainerStatus() // sidecar status listed before executor's container status + .withName("spark-executor-sidecar") + .withImage("k8s-spark-sidecar") + .withNewState() + .withNewTerminated() + .withMessage("Failed") + .withExitCode(2) + .endTerminated() + .endState() + .endContainerStatus() + .addNewContainerStatus() + .withName(TEST_SPARK_EXECUTOR_CONTAINER_NAME) + .withImage("k8s-spark") .withNewState() .withNewTerminated() .withMessage("Failed") @@ -200,7 +233,7 @@ object ExecutorLifecycleTestUtils { .endSpec() .build() val container = new ContainerBuilder() - .withName("spark-executor") + .withName(TEST_SPARK_EXECUTOR_CONTAINER_NAME) .withImage("k8s-spark") .build() SparkPod(pod, container) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala index 96be5dfabd121..d3b7213807afb 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala @@ -33,6 +33,7 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.{SparkConf, SparkFunSuite} import org.apache.spark.deploy.k8s.Config +import org.apache.spark.deploy.k8s.Config._ import org.apache.spark.deploy.k8s.Constants._ import org.apache.spark.deploy.k8s.Fabric8Aliases._ import org.apache.spark.deploy.k8s.KubernetesUtils._ @@ -60,6 +61,8 @@ class ExecutorPodsLifecycleManagerSuite extends SparkFunSuite with BeforeAndAfte before { MockitoAnnotations.openMocks(this).close() + val sparkConf = new SparkConf() + .set(KUBERNETES_EXECUTOR_PODTEMPLATE_CONTAINER_NAME, TEST_SPARK_EXECUTOR_CONTAINER_NAME) snapshotsStore = new DeterministicExecutorPodsSnapshotsStore() namedExecutorPods = mutable.Map.empty[String, PodResource] when(schedulerBackend.getExecutorsWithRegistrationTs()).thenReturn(Map.empty[String, Long]) @@ -67,7 +70,7 @@ class ExecutorPodsLifecycleManagerSuite extends SparkFunSuite with BeforeAndAfte when(podOperations.inNamespace(anyString())).thenReturn(podsWithNamespace) when(podsWithNamespace.withName(any(classOf[String]))).thenAnswer(namedPodsAnswer()) eventHandlerUnderTest = new ExecutorPodsLifecycleManager( - new SparkConf(), + sparkConf, kubernetesClient, snapshotsStore) eventHandlerUnderTest.start(schedulerBackend) @@ -162,6 +165,15 @@ class ExecutorPodsLifecycleManagerSuite extends SparkFunSuite with BeforeAndAfte .edit(any[UnaryOperator[Pod]]()) } + test("SPARK-49804: Use the exit code of executor container always") { + val failedPod = failedExecutorWithSidecarStatusListedFirst(1) + snapshotsStore.updatePod(failedPod) + snapshotsStore.notifySubscribers() + val msg = exitReasonMessage(1, failedPod, 1) + val expectedLossReason = ExecutorExited(1, exitCausedByApp = true, msg) + verify(schedulerBackend).doRemoveExecutor("1", expectedLossReason) + } + private def exitReasonMessage(execId: Int, failedPod: Pod, exitCode: Int): String = { val reason = Option(failedPod.getStatus.getReason) val message = Option(failedPod.getStatus.getMessage) From f18c4e7722b46e8573e959f5f3b063ed0efa5d23 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 27 Sep 2024 15:27:34 +0800 Subject: [PATCH 157/189] [SPARK-49805][SQL][ML] Remove private[xxx] functions from `function.scala` ### What changes were proposed in this pull request? Remove private[xxx] functions from `function.scala` ### Why are the changes needed? internal functions can be directly invoked by `Column.internalFn`, no need to add them in `function.scala` ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48276 from zhengruifeng/move_private_func. Authored-by: Ruifeng Zheng Signed-off-by: yangjie01 --- .../main/scala/org/apache/spark/ml/recommendation/ALS.scala | 5 ++++- .../apache/spark/ml/recommendation/CollectTopKSuite.scala | 3 ++- sql/api/src/main/scala/org/apache/spark/sql/functions.scala | 3 --- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala index 1a004f71749e1..5899bf891ec9d 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala @@ -517,7 +517,7 @@ class ALSModel private[ml] ( ) ratings.groupBy(srcOutputColumn) - .agg(collect_top_k(struct(ratingColumn, dstOutputColumn), num, false)) + .agg(ALSModel.collect_top_k(struct(ratingColumn, dstOutputColumn), num, false)) .as[(Int, Seq[(Float, Int)])] .map(t => (t._1, t._2.map(p => (p._2, p._1)))) .toDF(srcOutputColumn, recommendColumn) @@ -546,6 +546,9 @@ object ALSModel extends MLReadable[ALSModel] { private val Drop = "drop" private[recommendation] final val supportedColdStartStrategies = Array(NaN, Drop) + private[recommendation] def collect_top_k(e: Column, num: Int, reverse: Boolean): Column = + Column.internalFn("collect_top_k", e, lit(num), lit(reverse)) + @Since("1.6.0") override def read: MLReader[ALSModel] = new ALSModelReader diff --git a/mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectTopKSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectTopKSuite.scala index b79e10d0d267e..bd83d5498ae6f 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectTopKSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/recommendation/CollectTopKSuite.scala @@ -17,9 +17,10 @@ package org.apache.spark.ml.recommendation +import org.apache.spark.ml.recommendation.ALSModel.collect_top_k import org.apache.spark.ml.util.MLTest import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions.{col, collect_top_k, struct} +import org.apache.spark.sql.functions.{col, struct} class CollectTopKSuite extends MLTest { diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 93bff22621057..e6fd06f2ec632 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -401,9 +401,6 @@ object functions { def count_min_sketch(e: Column, eps: Column, confidence: Column): Column = count_min_sketch(e, eps, confidence, lit(SparkClassUtils.random.nextLong)) - private[spark] def collect_top_k(e: Column, num: Int, reverse: Boolean): Column = - Column.internalFn("collect_top_k", e, lit(num), lit(reverse)) - /** * Aggregate function: returns the Pearson Correlation Coefficient for two columns. * From 9b739d415cd51c8dd3f9332bae225196bab17d48 Mon Sep 17 00:00:00 2001 From: Mikhail Nikoliukin Date: Fri, 27 Sep 2024 21:48:35 +0800 Subject: [PATCH 158/189] [SPARK-49757][SQL] Support IDENTIFIER expression in SET CATALOG statement ### What changes were proposed in this pull request? This pr adds possibility to use `IDENTIFIER(...)` for a catalog name in `SET CATALOG` statement. For instance `SET CATALOG IDENTIFIER('test')` now works the same as `SET CATALOG test` ### Why are the changes needed? 1. Consistency of API. It can be confusing for user that he can use IDENTIFIER in some contexts but cannot for catalogs. 2. Parametrization. It allows user to write `SET CATALOG IDENTIFIER(:user_data)` and doesn't worry about SQL injections. ### Does this PR introduce _any_ user-facing change? Yes, now `SET CATALOG IDENTIFIER(...)` works. It can be used with any string expressions and parametrization. But multipart identifiers (like `IDENTIFIER('database.table')`) are banned and will rise ParseException with new type `INVALID_SQL_SYNTAX.MULTI_PART_CATALOG_NAME`. This restriction always has been on grammar level, but now user can try to bind such identifier via parameters. ### How was this patch tested? Unit tests with several new covering new behavior. ### Was this patch authored or co-authored using generative AI tooling? Yes, some code suggestions Generated-by: GitHub Copilot Closes #48228 from mikhailnik-db/SPARK-49757. Authored-by: Mikhail Nikoliukin Signed-off-by: Wenchen Fan --- .../resources/error/error-conditions.json | 2 +- .../sql/catalyst/parser/SqlBaseParser.g4 | 8 +++- .../spark/sql/errors/QueryParsingErrors.scala | 14 +++++-- .../spark/sql/execution/SparkSqlParser.scala | 35 ++++++++++++---- .../identifier-clause.sql.out | 2 +- .../results/identifier-clause.sql.out | 2 +- .../sql/connector/DataSourceV2SQLSuite.scala | 42 +++++++++++++++++++ .../sql/errors/QueryParsingErrorsSuite.scala | 4 +- .../execution/command/DDLParserSuite.scala | 4 +- 9 files changed, 93 insertions(+), 20 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index e83202d9e5ee3..3fcb53426eccf 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -3023,7 +3023,7 @@ }, "MULTI_PART_NAME" : { "message" : [ - " with multiple part function name() is not allowed." + " with multiple part name() is not allowed." ] }, "OPTION_IS_INVALID" : { diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 094f7f5315b80..866634b041280 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -148,7 +148,7 @@ statement | ctes? dmlStatementNoWith #dmlStatement | USE identifierReference #use | USE namespace identifierReference #useNamespace - | SET CATALOG (errorCapturingIdentifier | stringLit) #setCatalog + | SET CATALOG catalogIdentifierReference #setCatalog | CREATE namespace (IF errorCapturingNot EXISTS)? identifierReference (commentSpec | locationSpec | @@ -594,6 +594,12 @@ identifierReference | multipartIdentifier ; +catalogIdentifierReference + : IDENTIFIER_KW LEFT_PAREN expression RIGHT_PAREN + | errorCapturingIdentifier + | stringLit + ; + queryOrganization : (ORDER BY order+=sortItem (COMMA order+=sortItem)*)? (CLUSTER BY clusterBy+=expression (COMMA clusterBy+=expression)*)? diff --git a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala index b19607a28f06c..b0743d6de4772 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala @@ -621,9 +621,8 @@ private[sql] object QueryParsingErrors extends DataTypeErrorsBase { def unsupportedFunctionNameError(funcName: Seq[String], ctx: ParserRuleContext): Throwable = { new ParseException( errorClass = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", - messageParameters = Map( - "statement" -> toSQLStmt("CREATE TEMPORARY FUNCTION"), - "funcName" -> toSQLId(funcName)), + messageParameters = + Map("statement" -> toSQLStmt("CREATE TEMPORARY FUNCTION"), "name" -> toSQLId(funcName)), ctx) } @@ -665,7 +664,14 @@ private[sql] object QueryParsingErrors extends DataTypeErrorsBase { new ParseException( errorClass = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", messageParameters = - Map("statement" -> toSQLStmt("DROP TEMPORARY FUNCTION"), "funcName" -> toSQLId(name)), + Map("statement" -> toSQLStmt("DROP TEMPORARY FUNCTION"), "name" -> toSQLId(name)), + ctx) + } + + def invalidNameForSetCatalog(name: Seq[String], ctx: ParserRuleContext): Throwable = { + new ParseException( + errorClass = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", + messageParameters = Map("statement" -> toSQLStmt("SET CATALOG"), "name" -> toSQLId(name)), ctx) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala index a8261e5d98ba0..1c735154f25ed 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala @@ -27,7 +27,7 @@ import org.antlr.v4.runtime.tree.TerminalNode import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} -import org.apache.spark.sql.catalyst.analysis.{GlobalTempView, LocalTempView, PersistedView, SchemaEvolution, SchemaTypeEvolution, UnresolvedFunctionName, UnresolvedIdentifier, UnresolvedNamespace} +import org.apache.spark.sql.catalyst.analysis.{GlobalTempView, LocalTempView, PersistedView, PlanWithUnresolvedIdentifier, SchemaEvolution, SchemaTypeEvolution, UnresolvedFunctionName, UnresolvedIdentifier, UnresolvedNamespace} import org.apache.spark.sql.catalyst.catalog._ import org.apache.spark.sql.catalyst.expressions.{Expression, Literal} import org.apache.spark.sql.catalyst.parser._ @@ -67,6 +67,25 @@ class SparkSqlAstBuilder extends AstBuilder { private val configValueDef = """([^;]*);*""".r private val strLiteralDef = """(".*?[^\\]"|'.*?[^\\]'|[^ \n\r\t"']+)""".r + private def withCatalogIdentClause( + ctx: CatalogIdentifierReferenceContext, + builder: Seq[String] => LogicalPlan): LogicalPlan = { + val exprCtx = ctx.expression + if (exprCtx != null) { + // resolve later in analyzer + PlanWithUnresolvedIdentifier(withOrigin(exprCtx) { expression(exprCtx) }, Nil, + (ident, _) => builder(ident)) + } else if (ctx.errorCapturingIdentifier() != null) { + // resolve immediately + builder.apply(Seq(ctx.errorCapturingIdentifier().getText)) + } else if (ctx.stringLit() != null) { + // resolve immediately + builder.apply(Seq(string(visitStringLit(ctx.stringLit())))) + } else { + throw SparkException.internalError("Invalid catalog name") + } + } + /** * Create a [[SetCommand]] logical plan. * @@ -276,13 +295,13 @@ class SparkSqlAstBuilder extends AstBuilder { * Create a [[SetCatalogCommand]] logical command. */ override def visitSetCatalog(ctx: SetCatalogContext): LogicalPlan = withOrigin(ctx) { - if (ctx.errorCapturingIdentifier() != null) { - SetCatalogCommand(ctx.errorCapturingIdentifier().getText) - } else if (ctx.stringLit() != null) { - SetCatalogCommand(string(visitStringLit(ctx.stringLit()))) - } else { - throw SparkException.internalError("Invalid catalog name") - } + withCatalogIdentClause(ctx.catalogIdentifierReference, identifiers => { + if (identifiers.size > 1) { + // can occur when user put multipart string in IDENTIFIER(...) clause + throw QueryParsingErrors.invalidNameForSetCatalog(identifiers, ctx) + } + SetCatalogCommand(identifiers.head) + }) } /** diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out index f0bf8b883dd8b..20e6ca1e6a2ec 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out @@ -893,7 +893,7 @@ org.apache.spark.sql.catalyst.parser.ParseException "errorClass" : "INVALID_SQL_SYNTAX.MULTI_PART_NAME", "sqlState" : "42000", "messageParameters" : { - "funcName" : "`default`.`myDoubleAvg`", + "name" : "`default`.`myDoubleAvg`", "statement" : "DROP TEMPORARY FUNCTION" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out b/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out index 952fb8fdc2bd2..596745b4ba5d8 100644 --- a/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out @@ -1024,7 +1024,7 @@ org.apache.spark.sql.catalyst.parser.ParseException "errorClass" : "INVALID_SQL_SYNTAX.MULTI_PART_NAME", "sqlState" : "42000", "messageParameters" : { - "funcName" : "`default`.`myDoubleAvg`", + "name" : "`default`.`myDoubleAvg`", "statement" : "DROP TEMPORARY FUNCTION" }, "queryContext" : [ { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index dac066bbef838..6b58d23e92603 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -2887,6 +2887,48 @@ class DataSourceV2SQLSuiteV1Filter "config" -> "\"spark.sql.catalog.not_exist_catalog\"")) } + test("SPARK-49757: SET CATALOG statement with IDENTIFIER should work") { + val catalogManager = spark.sessionState.catalogManager + assert(catalogManager.currentCatalog.name() == SESSION_CATALOG_NAME) + + sql("SET CATALOG IDENTIFIER('testcat')") + assert(catalogManager.currentCatalog.name() == "testcat") + + spark.sql("SET CATALOG IDENTIFIER(:param)", Map("param" -> "testcat2")) + assert(catalogManager.currentCatalog.name() == "testcat2") + + checkError( + exception = intercept[CatalogNotFoundException] { + sql("SET CATALOG IDENTIFIER('not_exist_catalog')") + }, + condition = "CATALOG_NOT_FOUND", + parameters = Map( + "catalogName" -> "`not_exist_catalog`", + "config" -> "\"spark.sql.catalog.not_exist_catalog\"") + ) + } + + test("SPARK-49757: SET CATALOG statement with IDENTIFIER with multipart name should fail") { + val catalogManager = spark.sessionState.catalogManager + assert(catalogManager.currentCatalog.name() == SESSION_CATALOG_NAME) + + val sqlText = "SET CATALOG IDENTIFIER(:param)" + checkError( + exception = intercept[ParseException] { + spark.sql(sqlText, Map("param" -> "testcat.ns1")) + }, + condition = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", + parameters = Map( + "name" -> "`testcat`.`ns1`", + "statement" -> "SET CATALOG" + ), + context = ExpectedContext( + fragment = sqlText, + start = 0, + stop = 29) + ) + } + test("SPARK-35973: ShowCatalogs") { val schema = new StructType() .add("catalog", StringType, nullable = false) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryParsingErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryParsingErrorsSuite.scala index da7b6e7f63c85..666f85e19c1c3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryParsingErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryParsingErrorsSuite.scala @@ -334,7 +334,7 @@ class QueryParsingErrorsSuite extends QueryTest with SharedSparkSession with SQL sqlState = "42000", parameters = Map( "statement" -> "CREATE TEMPORARY FUNCTION", - "funcName" -> "`ns`.`db`.`func`"), + "name" -> "`ns`.`db`.`func`"), context = ExpectedContext( fragment = sqlText, start = 0, @@ -367,7 +367,7 @@ class QueryParsingErrorsSuite extends QueryTest with SharedSparkSession with SQL sqlState = "42000", parameters = Map( "statement" -> "DROP TEMPORARY FUNCTION", - "funcName" -> "`db`.`func`"), + "name" -> "`db`.`func`"), context = ExpectedContext( fragment = sqlText, start = 0, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala index 176eb7c290764..8b868c0e17230 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala @@ -688,7 +688,7 @@ class DDLParserSuite extends AnalysisTest with SharedSparkSession { checkError( exception = parseException(sql1), condition = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", - parameters = Map("statement" -> "DROP TEMPORARY FUNCTION", "funcName" -> "`a`.`b`"), + parameters = Map("statement" -> "DROP TEMPORARY FUNCTION", "name" -> "`a`.`b`"), context = ExpectedContext( fragment = sql1, start = 0, @@ -698,7 +698,7 @@ class DDLParserSuite extends AnalysisTest with SharedSparkSession { checkError( exception = parseException(sql2), condition = "INVALID_SQL_SYNTAX.MULTI_PART_NAME", - parameters = Map("statement" -> "DROP TEMPORARY FUNCTION", "funcName" -> "`a`.`b`"), + parameters = Map("statement" -> "DROP TEMPORARY FUNCTION", "name" -> "`a`.`b`"), context = ExpectedContext( fragment = sql2, start = 0, From d7abddc454ffef6ac16e8f6df6f601eec621ddfd Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 27 Sep 2024 21:53:10 +0800 Subject: [PATCH 159/189] [SPARK-49808][SQL] Fix a deadlock in subquery execution due to lazy vals ### What changes were proposed in this pull request? Fix a deadlock in subquery execution due to lazy vals ### Why are the changes needed? we observed a deadlock between `QueryPlan.canonicalized` and `QueryPlan.references`: ``` 24/09/04 04:46:54 ERROR DeadlockDetector: Found 2 new deadlock thread(s): "ScalaTest-run-running-SubquerySuite" prio=5 Id=1 BLOCKED on org.apache.spark.sql.execution.aggregate.HashAggregateExec87abc7f owned by "subquery-5" Id=112 at app//org.apache.spark.sql.catalyst.plans.QueryPlan.canonicalized$lzycompute(QueryPlan.scala:684) - blocked on org.apache.spark.sql.execution.aggregate.HashAggregateExec87abc7f at app//org.apache.spark.sql.catalyst.plans.QueryPlan.canonicalized(QueryPlan.scala:684) at app//org.apache.spark.sql.catalyst.plans.QueryPlan.$anonfun$doCanonicalize$2(QueryPlan.scala:716) at app//org.apache.spark.sql.catalyst.plans.QueryPlan$$Lambda$4058/0x00007f740f3d0cb0.apply(Unknown Source) at app//org.apache.spark.sql.catalyst.trees.UnaryLike.mapChildren(TreeNode.scala:1314) at app//org.apache.spark.sql.catalyst.trees.UnaryLike.mapChildren$(TreeNode.scala:1313) at app//org.apache.spark.sql.execution.WholeStageCodegenExec.mapChildren(WholeStageCodegenExec.scala:639) at app//org.apache.spark.sql.catalyst.plans.QueryPlan.doCanonicalize(QueryPlan.scala:716) ... "subquery-5" daemon prio=5 Id=112 BLOCKED on org.apache.spark.sql.execution.WholeStageCodegenExec132a3243 owned by "ScalaTest-run-running-SubquerySuite" Id=1 at app//org.apache.spark.sql.catalyst.plans.QueryPlan.references$lzycompute(QueryPlan.scala:101) - blocked on org.apache.spark.sql.execution.WholeStageCodegenExec132a3243 at app//org.apache.spark.sql.catalyst.plans.QueryPlan.references(QueryPlan.scala:101) at app//org.apache.spark.sql.execution.CodegenSupport.usedInputs(WholeStageCodegenExec.scala:325) at app//org.apache.spark.sql.execution.CodegenSupport.usedInputs$(WholeStageCodegenExec.scala:325) at app//org.apache.spark.sql.execution.WholeStageCodegenExec.usedInputs(WholeStageCodegenExec.scala:639) at app//org.apache.spark.sql.execution.CodegenSupport.consume(WholeStageCodegenExec.scala:187) at app//org.apache.spark.sql.execution.CodegenSupport.consume$(WholeStageCodegenExec.scala:157) at app//org.apache.spark.sql.execution.aggregate.HashAggregateExec.consume(HashAggregateExec.scala:53) ``` The main thread `TakeOrderedAndProject.doExecute` is trying to compute `outputOrdering`, it top-down traverse the tree, and requires the lock of `QueryPlan.canonicalized` in the path. In this deadlock, it successfully obtained the lock of `WholeStageCodegenExec` and requires the lock of `HashAggregateExec`; Concurrently, a subquery execution thread is performing code generation and bottom-up traverses the tree via `def consume`, which checks `WholeStageCodegenExec.usedInputs` and refererences a lazy val `QueryPlan.references`. It requires the lock of `QueryPlan.references` in the path. In this deadlock, it successfully obtained the lock of `HashAggregateExec` and requires the lock of `WholeStageCodegenExec`; This is due to Scala's lazy val internally calls this.synchronized on the instance that contains the val. This creates a potential for deadlocks. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? manually checked with `com.databricks.spark.sql.SubquerySuite` we encountered this issue multiple times before this fix in `SubquerySuite`, and after this fix we didn't hit this issue in multiple runs. ### Was this patch authored or co-authored using generative AI tooling? no Closes #48279 from zhengruifeng/fix_deadlock. Authored-by: Ruifeng Zheng Signed-off-by: Wenchen Fan --- .../org/apache/spark/sql/catalyst/plans/QueryPlan.scala | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala index 3f417644082c3..ca5ff78b10e91 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.trees.TreePatternBits import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.util.LazyTry import org.apache.spark.util.collection.BitSet /** @@ -94,9 +95,11 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] * All Attributes that appear in expressions from this operator. Note that this set does not * include attributes that are implicitly referenced by being passed through to the output tuple. */ + def references: AttributeSet = lazyReferences.get + @transient - lazy val references: AttributeSet = { - AttributeSet.fromAttributeSets(expressions.map(_.references)) -- producedAttributes + private val lazyReferences = LazyTry { + AttributeSet(expressions) -- producedAttributes } /** From dd692e90b7384a789142cfccff0dbf10cead6a21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B8rn=20J=C3=B8rgensen?= Date: Fri, 27 Sep 2024 07:49:06 -0700 Subject: [PATCH 160/189] [SPARK-49801][FOLLOWUP][INFRA] Update `pandas` to 2.2.3 in `pages.yml` too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Fix doc build. ### Why are the changes needed? in https://github.com/apache/spark/pull/48269 > Oh, this seems to break GitHub Action Jekyll. https://github.com/apache/spark/actions/runs/11063509911/job/30742286270 Traceback (most recent call last): File "/opt/hostedtoolcache/Python/3.9.20/x64/lib/python3.9/site-packages/sphinx/config.py", line 332, in eval_config_file exec(code, namespace) File "/home/runner/work/spark/spark/python/docs/source/conf.py", line 33, in generate_supported_api(output_rst_file_path) File "/home/runner/work/spark/spark/python/pyspark/pandas/supported_api_gen.py", line 102, in generate_supported_api _check_pandas_version() File "/home/runner/work/spark/spark/python/pyspark/pandas/supported_api_gen.py", line 116, in _check_pandas_version raise ImportError(msg) ImportError: Warning: pandas 2.2.3 is required; your version is 2.2.2" ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48278 from bjornjorgensen/fix-pandas-2.2.3. Authored-by: Bjørn Jørgensen Signed-off-by: Dongjoon Hyun --- .github/workflows/pages.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 8faeb0557fbfb..f78f7895a183f 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -60,7 +60,7 @@ jobs: - name: Install Python dependencies run: | pip install 'sphinx==4.5.0' mkdocs 'pydata_sphinx_theme>=0.13' sphinx-copybutton nbsphinx numpydoc jinja2 markupsafe 'pyzmq<24.0.0' \ - ipython ipython_genutils sphinx_plotly_directive 'numpy>=1.20.0' pyarrow 'pandas==2.2.2' 'plotly>=4.8' 'docutils<0.18.0' \ + ipython ipython_genutils sphinx_plotly_directive 'numpy>=1.20.0' pyarrow 'pandas==2.2.3' 'plotly>=4.8' 'docutils<0.18.0' \ 'flake8==3.9.0' 'mypy==1.8.0' 'pytest==7.1.3' 'pytest-mypy-plugins==1.9.3' 'black==23.9.1' \ 'pandas-stubs==1.2.0.53' 'grpcio==1.62.0' 'grpcio-status==1.62.0' 'grpc-stubs==1.24.11' 'googleapis-common-protos-stubs==2.2.0' \ 'sphinxcontrib-applehelp==1.0.4' 'sphinxcontrib-devhelp==1.0.2' 'sphinxcontrib-htmlhelp==2.0.1' 'sphinxcontrib-qthelp==1.0.3' 'sphinxcontrib-serializinghtml==1.1.5' From 6dc628c31cdf48769ccd80cd2b81f7bd6386276f Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 27 Sep 2024 08:51:47 -0700 Subject: [PATCH 161/189] [SPARK-49809][BUILD] Use `sbt.IO` in `SparkBuild.scala` to avoid naming conflicts with `java.io.IO` in Java 23 ### What changes were proposed in this pull request? This pr change to use `sbt.IO` in `SparkBuild.scala` to avoid naming conflicts with `java.io.IO` in Java 23, and after this PR, Spark can be built using sbt with Java 23(current pr does not focus on the results of `sbt/test` with Java 23) ### Why are the changes needed? Make Spark be compiled using sbt with Java 23. Because Java 23 has added `java.io.IO`, and `SparkBuild.scala` imports both `java.io._` and `sbt._`, this results in the following error when executing ``` build/sbt -Phadoop-3 -Phive-thriftserver -Pspark-ganglia-lgpl -Pdocker-integration-tests -Pyarn -Pvolcano -Pkubernetes -Pkinesis-asl -Phive -Phadoop-cloud Test/package streaming-kinesis-asl-assembly/assembly connect/assembly ``` with Java 23 ``` build/sbt -Phadoop-3 -Phive-thriftserver -Pspark-ganglia-lgpl -Pdocker-integration-tests -Pyarn -Pvolcano -Pkubernetes -Pkinesis-asl -Phive -Phadoop-cloud Test/package streaming-kinesis-asl-assembly/assembly connect/assembly Using /Users/yangjie01/Tools/zulu23 as default JAVA_HOME. Note, this will be overridden by -java-home if it is set. [info] welcome to sbt 1.9.3 (Azul Systems, Inc. Java 23) [info] loading settings for project global-plugins from idea.sbt ... [info] loading global plugins from /Users/yangjie01/.sbt/1.0/plugins [info] loading settings for project spark-sbt-build from plugins.sbt ... [info] loading project definition from /Users/yangjie01/SourceCode/git/spark-sbt/project [info] compiling 3 Scala sources to /Users/yangjie01/SourceCode/git/spark-sbt/project/target/scala-2.12/sbt-1.0/classes ... [error] /Users/yangjie01/SourceCode/git/spark-sbt/project/SparkBuild.scala:1209:7: reference to IO is ambiguous; [error] it is imported twice in the same scope by [error] import sbt._ [error] and import java.io._ [error] IO.write(file, s"$hadoopProvidedProp = $isHadoopProvided") [error] ^ [error] one error found [error] (Compile / compileIncremental) Compilation failed ``` ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? - Pass Github Actions - Manual check: ``` build/sbt -Phadoop-3 -Phive-thriftserver -Pspark-ganglia-lgpl -Pdocker-integration-tests -Pyarn -Pvolcano -Pkubernetes -Pkinesis-asl -Phive -Phadoop-cloud Test/package streaming-kinesis-asl-assembly/assembly connect/assembly ``` with Java 23, after this pr, the aforementioned command can be executed successfully. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48280 from LuciferYang/build-with-java23. Authored-by: yangjie01 Signed-off-by: Dongjoon Hyun --- project/SparkBuild.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index 82950fb30287a..6137984a53c0a 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -1206,7 +1206,7 @@ object YARN { genConfigProperties := { val file = (Compile / classDirectory).value / s"org/apache/spark/deploy/yarn/$propFileName" val isHadoopProvided = SbtPomKeys.effectivePom.value.getProperties.get(hadoopProvidedProp) - IO.write(file, s"$hadoopProvidedProp = $isHadoopProvided") + sbt.IO.write(file, s"$hadoopProvidedProp = $isHadoopProvided") }, Compile / copyResources := (Def.taskDyn { val c = (Compile / copyResources).value From b6681fbf32fa3596d7649d413f20cc5c6da64991 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 27 Sep 2024 13:42:54 -0700 Subject: [PATCH 162/189] [SPARK-49787][SQL] Cast between UDT and other types ### What changes were proposed in this pull request? This patch adds UDT support to `Cast` expression. ### Why are the changes needed? Our customer faced an error when migrating queries that write UDT column from Hive to Iceberg table. The error happens when Spark tries to cast UDT column to the data type (i.e., the sql type of the UDT) of the table column. The cast is added by table column resolution rule for V2 writing commands. Currently `Cast` expression doesn't support casting between UDT and other types. However, underlying an UDT, it is serialized as its `sqlType`, `Cast` should be able to cast between the `sqlType` and other types. ### Does this PR introduce _any_ user-facing change? Yes. User query can cast between UDT and other types. ### How was this patch tested? Unit test ### Was this patch authored or co-authored using generative AI tooling? No Closes #48251 from viirya/cast_udt. Authored-by: Liang-Chi Hsieh Signed-off-by: huaxingao --- python/pyspark/sql/tests/test_types.py | 16 +- .../apache/spark/sql/types/UpCastRule.scala | 4 + .../spark/sql/catalyst/expressions/Cast.scala | 175 ++++++++++-------- .../sql/catalyst/expressions/literals.scala | 84 +++++---- .../catalyst/expressions/CastSuiteBase.scala | 42 ++++- 5 files changed, 202 insertions(+), 119 deletions(-) diff --git a/python/pyspark/sql/tests/test_types.py b/python/pyspark/sql/tests/test_types.py index 8610ace52d86a..c240a84d1edb9 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -28,7 +28,6 @@ from pyspark.sql import Row from pyspark.sql import functions as F from pyspark.errors import ( - AnalysisException, ParseException, PySparkTypeError, PySparkValueError, @@ -1130,10 +1129,17 @@ def test_cast_to_string_with_udt(self): def test_cast_to_udt_with_udt(self): row = Row(point=ExamplePoint(1.0, 2.0), python_only_point=PythonOnlyPoint(1.0, 2.0)) df = self.spark.createDataFrame([row]) - with self.assertRaises(AnalysisException): - df.select(F.col("point").cast(PythonOnlyUDT())).collect() - with self.assertRaises(AnalysisException): - df.select(F.col("python_only_point").cast(ExamplePointUDT())).collect() + result = df.select(F.col("point").cast(PythonOnlyUDT())).collect() + self.assertEqual( + result, + [Row(point=PythonOnlyPoint(1.0, 2.0))], + ) + + result = df.select(F.col("python_only_point").cast(ExamplePointUDT())).collect() + self.assertEqual( + result, + [Row(python_only_point=ExamplePoint(1.0, 2.0))], + ) def test_struct_type(self): struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/UpCastRule.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/UpCastRule.scala index 4993e249b3059..6f2fd41f1f799 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/UpCastRule.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/UpCastRule.scala @@ -66,6 +66,10 @@ private[sql] object UpCastRule { case (from: UserDefinedType[_], to: UserDefinedType[_]) if to.acceptsType(from) => true + case (udt: UserDefinedType[_], toType) => canUpCast(udt.sqlType, toType) + + case (fromType, udt: UserDefinedType[_]) => canUpCast(fromType, udt.sqlType) + case _ => false } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala index 7a2799e99fe2d..9a29cb4a2bfb3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala @@ -150,6 +150,10 @@ object Cast extends QueryErrorsBase { case (udt1: UserDefinedType[_], udt2: UserDefinedType[_]) if udt2.acceptsType(udt1) => true + case (udt: UserDefinedType[_], toType) => canAnsiCast(udt.sqlType, toType) + + case (fromType, udt: UserDefinedType[_]) => canAnsiCast(fromType, udt.sqlType) + case _ => false } @@ -267,6 +271,10 @@ object Cast extends QueryErrorsBase { case (udt1: UserDefinedType[_], udt2: UserDefinedType[_]) if udt2.acceptsType(udt1) => true + case (udt: UserDefinedType[_], toType) => canCast(udt.sqlType, toType) + + case (fromType, udt: UserDefinedType[_]) => canCast(fromType, udt.sqlType) + case _ => false } @@ -1123,33 +1131,42 @@ case class Cast( variant.VariantGet.cast(v, to, evalMode != EvalMode.TRY, timeZoneId, zoneId) }) } else { - to match { - case dt if dt == from => identity[Any] - case VariantType => input => variant.VariantExpressionEvalUtils.castToVariant(input, from) - case _: StringType => castToString(from) - case BinaryType => castToBinary(from) - case DateType => castToDate(from) - case decimal: DecimalType => castToDecimal(from, decimal) - case TimestampType => castToTimestamp(from) - case TimestampNTZType => castToTimestampNTZ(from) - case CalendarIntervalType => castToInterval(from) - case it: DayTimeIntervalType => castToDayTimeInterval(from, it) - case it: YearMonthIntervalType => castToYearMonthInterval(from, it) - case BooleanType => castToBoolean(from) - case ByteType => castToByte(from) - case ShortType => castToShort(from) - case IntegerType => castToInt(from) - case FloatType => castToFloat(from) - case LongType => castToLong(from) - case DoubleType => castToDouble(from) - case array: ArrayType => - castArray(from.asInstanceOf[ArrayType].elementType, array.elementType) - case map: MapType => castMap(from.asInstanceOf[MapType], map) - case struct: StructType => castStruct(from.asInstanceOf[StructType], struct) - case udt: UserDefinedType[_] if udt.acceptsType(from) => - identity[Any] - case _: UserDefinedType[_] => - throw QueryExecutionErrors.cannotCastError(from, to) + from match { + // `castToString` has special handling for `UserDefinedType` + case udt: UserDefinedType[_] if !to.isInstanceOf[StringType] => + castInternal(udt.sqlType, to) + case _ => + to match { + case dt if dt == from => identity[Any] + case VariantType => input => + variant.VariantExpressionEvalUtils.castToVariant(input, from) + case _: StringType => castToString(from) + case BinaryType => castToBinary(from) + case DateType => castToDate(from) + case decimal: DecimalType => castToDecimal(from, decimal) + case TimestampType => castToTimestamp(from) + case TimestampNTZType => castToTimestampNTZ(from) + case CalendarIntervalType => castToInterval(from) + case it: DayTimeIntervalType => castToDayTimeInterval(from, it) + case it: YearMonthIntervalType => castToYearMonthInterval(from, it) + case BooleanType => castToBoolean(from) + case ByteType => castToByte(from) + case ShortType => castToShort(from) + case IntegerType => castToInt(from) + case FloatType => castToFloat(from) + case LongType => castToLong(from) + case DoubleType => castToDouble(from) + case array: ArrayType => + castArray(from.asInstanceOf[ArrayType].elementType, array.elementType) + case map: MapType => castMap(from.asInstanceOf[MapType], map) + case struct: StructType => castStruct(from.asInstanceOf[StructType], struct) + case udt: UserDefinedType[_] if udt.acceptsType(from) => + identity[Any] + case udt: UserDefinedType[_] => + castInternal(from, udt.sqlType) + case _ => + throw QueryExecutionErrors.cannotCastError(from, to) + } } } } @@ -1211,54 +1228,64 @@ case class Cast( private[this] def nullSafeCastFunction( from: DataType, to: DataType, - ctx: CodegenContext): CastFunction = to match { - - case _ if from == NullType => (c, evPrim, evNull) => code"$evNull = true;" - case _ if to == from => (c, evPrim, evNull) => code"$evPrim = $c;" - case _ if from.isInstanceOf[VariantType] => (c, evPrim, evNull) => - val tmp = ctx.freshVariable("tmp", classOf[Object]) - val dataTypeArg = ctx.addReferenceObj("dataType", to) - val zoneStrArg = ctx.addReferenceObj("zoneStr", timeZoneId) - val zoneIdArg = ctx.addReferenceObj("zoneId", zoneId, classOf[ZoneId].getName) - val failOnError = evalMode != EvalMode.TRY - val cls = classOf[variant.VariantGet].getName - code""" - Object $tmp = $cls.cast($c, $dataTypeArg, $failOnError, $zoneStrArg, $zoneIdArg); - if ($tmp == null) { - $evNull = true; - } else { - $evPrim = (${CodeGenerator.boxedType(to)})$tmp; + ctx: CodegenContext): CastFunction = { + from match { + // `castToStringCode` has special handling for `UserDefinedType` + case udt: UserDefinedType[_] if !to.isInstanceOf[StringType] => + nullSafeCastFunction(udt.sqlType, to, ctx) + case _ => + to match { + + case _ if from == NullType => (c, evPrim, evNull) => code"$evNull = true;" + case _ if to == from => (c, evPrim, evNull) => code"$evPrim = $c;" + case _ if from.isInstanceOf[VariantType] => (c, evPrim, evNull) => + val tmp = ctx.freshVariable("tmp", classOf[Object]) + val dataTypeArg = ctx.addReferenceObj("dataType", to) + val zoneStrArg = ctx.addReferenceObj("zoneStr", timeZoneId) + val zoneIdArg = ctx.addReferenceObj("zoneId", zoneId, classOf[ZoneId].getName) + val failOnError = evalMode != EvalMode.TRY + val cls = classOf[variant.VariantGet].getName + code""" + Object $tmp = $cls.cast($c, $dataTypeArg, $failOnError, $zoneStrArg, $zoneIdArg); + if ($tmp == null) { + $evNull = true; + } else { + $evPrim = (${CodeGenerator.boxedType(to)})$tmp; + } + """ + case VariantType => + val cls = variant.VariantExpressionEvalUtils.getClass.getName.stripSuffix("$") + val fromArg = ctx.addReferenceObj("from", from) + (c, evPrim, evNull) => code"$evPrim = $cls.castToVariant($c, $fromArg);" + case _: StringType => (c, evPrim, _) => castToStringCode(from, ctx).apply(c, evPrim) + case BinaryType => castToBinaryCode(from) + case DateType => castToDateCode(from, ctx) + case decimal: DecimalType => castToDecimalCode(from, decimal, ctx) + case TimestampType => castToTimestampCode(from, ctx) + case TimestampNTZType => castToTimestampNTZCode(from, ctx) + case CalendarIntervalType => castToIntervalCode(from) + case it: DayTimeIntervalType => castToDayTimeIntervalCode(from, it) + case it: YearMonthIntervalType => castToYearMonthIntervalCode(from, it) + case BooleanType => castToBooleanCode(from, ctx) + case ByteType => castToByteCode(from, ctx) + case ShortType => castToShortCode(from, ctx) + case IntegerType => castToIntCode(from, ctx) + case FloatType => castToFloatCode(from, ctx) + case LongType => castToLongCode(from, ctx) + case DoubleType => castToDoubleCode(from, ctx) + + case array: ArrayType => + castArrayCode(from.asInstanceOf[ArrayType].elementType, array.elementType, ctx) + case map: MapType => castMapCode(from.asInstanceOf[MapType], map, ctx) + case struct: StructType => castStructCode(from.asInstanceOf[StructType], struct, ctx) + case udt: UserDefinedType[_] if udt.acceptsType(from) => + (c, evPrim, evNull) => code"$evPrim = $c;" + case udt: UserDefinedType[_] => + nullSafeCastFunction(from, udt.sqlType, ctx) + case _ => + throw QueryExecutionErrors.cannotCastError(from, to) } - """ - case VariantType => - val cls = variant.VariantExpressionEvalUtils.getClass.getName.stripSuffix("$") - val fromArg = ctx.addReferenceObj("from", from) - (c, evPrim, evNull) => code"$evPrim = $cls.castToVariant($c, $fromArg);" - case _: StringType => (c, evPrim, _) => castToStringCode(from, ctx).apply(c, evPrim) - case BinaryType => castToBinaryCode(from) - case DateType => castToDateCode(from, ctx) - case decimal: DecimalType => castToDecimalCode(from, decimal, ctx) - case TimestampType => castToTimestampCode(from, ctx) - case TimestampNTZType => castToTimestampNTZCode(from, ctx) - case CalendarIntervalType => castToIntervalCode(from) - case it: DayTimeIntervalType => castToDayTimeIntervalCode(from, it) - case it: YearMonthIntervalType => castToYearMonthIntervalCode(from, it) - case BooleanType => castToBooleanCode(from, ctx) - case ByteType => castToByteCode(from, ctx) - case ShortType => castToShortCode(from, ctx) - case IntegerType => castToIntCode(from, ctx) - case FloatType => castToFloatCode(from, ctx) - case LongType => castToLongCode(from, ctx) - case DoubleType => castToDoubleCode(from, ctx) - - case array: ArrayType => - castArrayCode(from.asInstanceOf[ArrayType].elementType, array.elementType, ctx) - case map: MapType => castMapCode(from.asInstanceOf[MapType], map, ctx) - case struct: StructType => castStructCode(from.asInstanceOf[StructType], struct, ctx) - case udt: UserDefinedType[_] if udt.acceptsType(from) => - (c, evPrim, evNull) => code"$evPrim = $c;" - case _: UserDefinedType[_] => - throw QueryExecutionErrors.cannotCastError(from, to) + } } // Since we need to cast input expressions recursively inside ComplexTypes, such as Map's diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala index 4cffc7f0b53a3..362bb9af1661e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala @@ -441,47 +441,53 @@ case class Literal (value: Any, dataType: DataType) extends LeafExpression { override def eval(input: InternalRow): Any = value override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { - val javaType = CodeGenerator.javaType(dataType) - if (value == null) { - ExprCode.forNullValue(dataType) - } else { - def toExprCode(code: String): ExprCode = { - ExprCode.forNonNullValue(JavaCode.literal(code, dataType)) - } - dataType match { - case BooleanType | IntegerType | DateType | _: YearMonthIntervalType => - toExprCode(value.toString) - case FloatType => - value.asInstanceOf[Float] match { - case v if v.isNaN => - toExprCode("Float.NaN") - case Float.PositiveInfinity => - toExprCode("Float.POSITIVE_INFINITY") - case Float.NegativeInfinity => - toExprCode("Float.NEGATIVE_INFINITY") - case _ => - toExprCode(s"${value}F") - } - case DoubleType => - value.asInstanceOf[Double] match { - case v if v.isNaN => - toExprCode("Double.NaN") - case Double.PositiveInfinity => - toExprCode("Double.POSITIVE_INFINITY") - case Double.NegativeInfinity => - toExprCode("Double.NEGATIVE_INFINITY") - case _ => - toExprCode(s"${value}D") - } - case ByteType | ShortType => - ExprCode.forNonNullValue(JavaCode.expression(s"($javaType)$value", dataType)) - case TimestampType | TimestampNTZType | LongType | _: DayTimeIntervalType => - toExprCode(s"${value}L") - case _ => - val constRef = ctx.addReferenceObj("literal", value, javaType) - ExprCode.forNonNullValue(JavaCode.global(constRef, dataType)) + def gen(ctx: CodegenContext, ev: ExprCode, dataType: DataType): ExprCode = { + val javaType = CodeGenerator.javaType(dataType) + if (value == null) { + ExprCode.forNullValue(dataType) + } else { + def toExprCode(code: String): ExprCode = { + ExprCode.forNonNullValue(JavaCode.literal(code, dataType)) + } + + dataType match { + case BooleanType | IntegerType | DateType | _: YearMonthIntervalType => + toExprCode(value.toString) + case FloatType => + value.asInstanceOf[Float] match { + case v if v.isNaN => + toExprCode("Float.NaN") + case Float.PositiveInfinity => + toExprCode("Float.POSITIVE_INFINITY") + case Float.NegativeInfinity => + toExprCode("Float.NEGATIVE_INFINITY") + case _ => + toExprCode(s"${value}F") + } + case DoubleType => + value.asInstanceOf[Double] match { + case v if v.isNaN => + toExprCode("Double.NaN") + case Double.PositiveInfinity => + toExprCode("Double.POSITIVE_INFINITY") + case Double.NegativeInfinity => + toExprCode("Double.NEGATIVE_INFINITY") + case _ => + toExprCode(s"${value}D") + } + case ByteType | ShortType => + ExprCode.forNonNullValue(JavaCode.expression(s"($javaType)$value", dataType)) + case TimestampType | TimestampNTZType | LongType | _: DayTimeIntervalType => + toExprCode(s"${value}L") + case udt: UserDefinedType[_] => + gen(ctx, ev, udt.sqlType) + case _ => + val constRef = ctx.addReferenceObj("literal", value, javaType) + ExprCode.forNonNullValue(JavaCode.global(constRef, dataType)) + } } } + gen(ctx, ev, dataType) } override def sql: String = (value, dataType) match { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuiteBase.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuiteBase.scala index e87b54339821f..f915d6efeb827 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuiteBase.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuiteBase.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.catalyst.expressions import java.sql.{Date, Timestamp} -import java.time.{Duration, LocalDate, LocalDateTime, Period} +import java.time.{Duration, LocalDate, LocalDateTime, Period, Year => JYear} import java.time.temporal.ChronoUnit import java.util.{Calendar, Locale, TimeZone} @@ -37,6 +37,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.sql.types.DataTypeTestUtils.{dayTimeIntervalTypes, yearMonthIntervalTypes} import org.apache.spark.sql.types.DayTimeIntervalType.{DAY, HOUR, MINUTE, SECOND} +import org.apache.spark.sql.types.TestUDT._ import org.apache.spark.sql.types.UpCastRule.numericPrecedence import org.apache.spark.sql.types.YearMonthIntervalType.{MONTH, YEAR} import org.apache.spark.unsafe.types.UTF8String @@ -1409,4 +1410,43 @@ abstract class CastSuiteBase extends SparkFunSuite with ExpressionEvalHelper { assert(!Cast(timestampLiteral, TimestampNTZType).resolved) assert(!Cast(timestampNTZLiteral, TimestampType).resolved) } + + test("SPARK-49787: Cast between UDT and other types") { + val value = new MyDenseVector(Array(1.0, 2.0, -1.0)) + val udtType = new MyDenseVectorUDT() + val targetType = ArrayType(DoubleType, containsNull = false) + + val serialized = udtType.serialize(value) + + checkEvaluation(Cast(new Literal(serialized, udtType), targetType), serialized) + checkEvaluation(Cast(new Literal(serialized, targetType), udtType), serialized) + + val year = JYear.parse("2024") + val yearUDTType = new YearUDT() + + val yearSerialized = yearUDTType.serialize(year) + + checkEvaluation(Cast(new Literal(yearSerialized, yearUDTType), IntegerType), 2024) + checkEvaluation(Cast(new Literal(2024, IntegerType), yearUDTType), yearSerialized) + + val yearString = UTF8String.fromString("2024") + checkEvaluation(Cast(new Literal(yearSerialized, yearUDTType), StringType), yearString) + checkEvaluation(Cast(new Literal(yearString, StringType), yearUDTType), yearSerialized) + } +} + +private[sql] class YearUDT extends UserDefinedType[JYear] { + override def sqlType: DataType = IntegerType + + override def serialize(obj: JYear): Int = { + obj.getValue + } + + def deserialize(datum: Any): JYear = datum match { + case value: Int => JYear.of(value) + } + + override def userClass: Class[JYear] = classOf[JYear] + + private[spark] override def asNullable: YearUDT = this } From 4d70954b1aeb10767cea82250eb975e2c85f1f3b Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Fri, 27 Sep 2024 14:21:33 -0700 Subject: [PATCH 163/189] [SPARK-49817][BUILD] Upgrade `gcs-connector` to `2.2.25` ### What changes were proposed in this pull request? This PR aims to upgrade `gcs-connector` to 2.2.25. ### Why are the changes needed? To bring the latest bug fixes. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs and manual test. ``` $ dev/make-distribution.sh -Phadoop-cloud $ cd dist $ export KEYFILE=~/.ssh/apache-spark.json $ export EMAIL=$(jq -r '.client_email' < $KEYFILE) $ export PRIVATE_KEY_ID=$(jq -r '.private_key_id' < $KEYFILE) $ export PRIVATE_KEY="$(jq -r '.private_key' < $KEYFILE)" $ bin/spark-shell \ -c spark.hadoop.fs.gs.auth.service.account.email=$EMAIL \ -c spark.hadoop.fs.gs.auth.service.account.private.key.id=$PRIVATE_KEY_ID \ -c spark.hadoop.fs.gs.auth.service.account.private.key="$PRIVATE_KEY" WARNING: Using incubator modules: jdk.incubator.vector Using Spark's default log4j profile: org/apache/spark/log4j2-pattern-layout-defaults.properties Setting default log level to "WARN". To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel). Welcome to ____ __ / __/__ ___ _____/ /__ _\ \/ _ \/ _ `/ __/ '_/ /___/ .__/\_,_/_/ /_/\_\ version 4.0.0-SNAPSHOT /_/ Using Scala version 2.13.15 (OpenJDK 64-Bit Server VM, Java 21.0.4) Type in expressions to have them evaluated. Type :help for more information. 24/09/27 09:34:53 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable Spark context Web UI available at http://localhost:4040 Spark context available as 'sc' (master = local[*], app id = local-1727454893738). Spark session available as 'spark'. scala> spark.read.text("gs://apache-spark-bucket/README.md").count() val res0: Long = 124 scala> spark.read.orc("examples/src/main/resources/users.orc").write.mode("overwrite").orc("gs://apache-spark-bucket/users.orc") scala> spark.read.orc("gs://apache-spark-bucket/users.orc").show() +------+--------------+----------------+ | name|favorite_color|favorite_numbers| +------+--------------+----------------+ |Alyssa| NULL| [3, 9, 15, 20]| | Ben| red| []| +------+--------------+----------------+ scala> ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48285 from dongjoon-hyun/SPARK-49817. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 2 +- pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index c9a32757554be..95a667ccfc72d 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -67,7 +67,7 @@ error_prone_annotations/2.26.1//error_prone_annotations-2.26.1.jar esdk-obs-java/3.20.4.2//esdk-obs-java-3.20.4.2.jar failureaccess/1.0.2//failureaccess-1.0.2.jar flatbuffers-java/24.3.25//flatbuffers-java-24.3.25.jar -gcs-connector/hadoop3-2.2.21/shaded/gcs-connector-hadoop3-2.2.21-shaded.jar +gcs-connector/hadoop3-2.2.25/shaded/gcs-connector-hadoop3-2.2.25-shaded.jar gmetric4j/1.0.10//gmetric4j-1.0.10.jar gson/2.11.0//gson-2.11.0.jar guava/33.2.1-jre//guava-33.2.1-jre.jar diff --git a/pom.xml b/pom.xml index 22048b55da27f..4bdb92d86a727 100644 --- a/pom.xml +++ b/pom.xml @@ -161,7 +161,7 @@ 0.12.8 - hadoop3-2.2.21 + hadoop3-2.2.25 4.5.14 4.4.16 From d813f5467e930ed4a22d2ea6aa4333cf379ea7f9 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Fri, 27 Sep 2024 19:41:37 -0400 Subject: [PATCH 164/189] [SPARK-49417][CONNECT][SQL] Add Shared StreamingQueryManager interface ### What changes were proposed in this pull request? This PR adds a shared StreamingQueryManager interface. ### Why are the changes needed? We are working on a shared Scala SQL interface for Classic and Connect. This change is part of this work. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48217 from hvanhovell/SPARK-49417. Authored-by: Herman van Hovell Signed-off-by: Herman van Hovell --- .../org/apache/spark/sql/SparkSession.scala | 1 + .../sql/streaming/StreamingQueryManager.scala | 93 ++----------- .../apache/spark/sql/api/SparkSession.scala | 11 +- .../spark/sql/api/StreamingQueryManager.scala | 130 ++++++++++++++++++ .../org/apache/spark/sql/SparkSession.scala | 7 +- .../sql/streaming/StreamingQueryManager.scala | 95 ++----------- 6 files changed, 169 insertions(+), 168 deletions(-) create mode 100644 sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQueryManager.scala diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index 1b41566ca1d1d..b31670c1da57e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -212,6 +212,7 @@ class SparkSession private[sql] ( /** @inheritdoc */ def readStream: DataStreamReader = new DataStreamReader(this) + /** @inheritdoc */ lazy val streams: StreamingQueryManager = new StreamingQueryManager(this) /** @inheritdoc */ diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala index 7efced227d6d1..647d29c714dbb 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala @@ -27,7 +27,7 @@ import org.apache.spark.connect.proto.Command import org.apache.spark.connect.proto.StreamingQueryManagerCommand import org.apache.spark.connect.proto.StreamingQueryManagerCommandResult import org.apache.spark.internal.Logging -import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.{api, SparkSession} import org.apache.spark.sql.connect.common.InvalidPlanInput /** @@ -36,7 +36,9 @@ import org.apache.spark.sql.connect.common.InvalidPlanInput * @since 3.5.0 */ @Evolving -class StreamingQueryManager private[sql] (sparkSession: SparkSession) extends Logging { +class StreamingQueryManager private[sql] (sparkSession: SparkSession) + extends api.StreamingQueryManager + with Logging { // Mapping from id to StreamingQueryListener. There's another mapping from id to // StreamingQueryListener on server side. This is used by removeListener() to find the id @@ -53,29 +55,17 @@ class StreamingQueryManager private[sql] (sparkSession: SparkSession) extends Lo streamingQueryListenerBus.close() } - /** - * Returns a list of active queries associated with this SQLContext - * - * @since 3.5.0 - */ + /** @inheritdoc */ def active: Array[StreamingQuery] = { executeManagerCmd(_.setActive(true)).getActive.getActiveQueriesList.asScala.map { q => RemoteStreamingQuery.fromStreamingQueryInstanceResponse(sparkSession, q) }.toArray } - /** - * Returns the query if there is an active query with the given id, or null. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def get(id: UUID): StreamingQuery = get(id.toString) - /** - * Returns the query if there is an active query with the given id, or null. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def get(id: String): StreamingQuery = { val response = executeManagerCmd(_.setGetQuery(id)) if (response.hasQuery) { @@ -85,52 +75,13 @@ class StreamingQueryManager private[sql] (sparkSession: SparkSession) extends Lo } } - /** - * Wait until any of the queries on the associated SQLContext has terminated since the creation - * of the context, or since `resetTerminated()` was called. If any query was terminated with an - * exception, then the exception will be thrown. - * - * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either - * return immediately (if the query was terminated by `query.stop()`), or throw the exception - * immediately (if the query was terminated with exception). Use `resetTerminated()` to clear - * past terminations and wait for new terminations. - * - * In the case where multiple queries have terminated since `resetTermination()` was called, if - * any query has terminated with exception, then `awaitAnyTermination()` will throw any of the - * exception. For correctly documenting exceptions across multiple queries, users need to stop - * all of them after any of them terminates with exception, and then check the - * `query.exception()` for each query. - * - * @throws StreamingQueryException - * if any query has terminated with an exception - * @since 3.5.0 - */ + /** @inheritdoc */ @throws[StreamingQueryException] def awaitAnyTermination(): Unit = { executeManagerCmd(_.getAwaitAnyTerminationBuilder.build()) } - /** - * Wait until any of the queries on the associated SQLContext has terminated since the creation - * of the context, or since `resetTerminated()` was called. Returns whether any query has - * terminated or not (multiple may have terminated). If any query has terminated with an - * exception, then the exception will be thrown. - * - * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either - * return `true` immediately (if the query was terminated by `query.stop()`), or throw the - * exception immediately (if the query was terminated with exception). Use `resetTerminated()` - * to clear past terminations and wait for new terminations. - * - * In the case where multiple queries have terminated since `resetTermination()` was called, if - * any query has terminated with exception, then `awaitAnyTermination()` will throw any of the - * exception. For correctly documenting exceptions across multiple queries, users need to stop - * all of them after any of them terminates with exception, and then check the - * `query.exception()` for each query. - * - * @throws StreamingQueryException - * if any query has terminated with an exception - * @since 3.5.0 - */ + /** @inheritdoc */ @throws[StreamingQueryException] def awaitAnyTermination(timeoutMs: Long): Boolean = { require(timeoutMs > 0, "Timeout has to be positive") @@ -139,40 +90,22 @@ class StreamingQueryManager private[sql] (sparkSession: SparkSession) extends Lo timeoutMs)).getAwaitAnyTermination.getTerminated } - /** - * Forget about past terminated queries so that `awaitAnyTermination()` can be used again to - * wait for new terminations. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def resetTerminated(): Unit = { executeManagerCmd(_.setResetTerminated(true)) } - /** - * Register a [[StreamingQueryListener]] to receive up-calls for life cycle events of - * [[StreamingQuery]]. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def addListener(listener: StreamingQueryListener): Unit = { streamingQueryListenerBus.append(listener) } - /** - * Deregister a [[StreamingQueryListener]]. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def removeListener(listener: StreamingQueryListener): Unit = { streamingQueryListenerBus.remove(listener) } - /** - * List all [[StreamingQueryListener]]s attached to this [[StreamingQueryManager]]. - * - * @since 3.5.0 - */ + /** @inheritdoc */ def listListeners(): Array[StreamingQueryListener] = { streamingQueryListenerBus.list() } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala index 0f73a94c3c4a4..4dfeb87a11d92 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/SparkSession.scala @@ -25,7 +25,7 @@ import _root_.java.lang import _root_.java.net.URI import _root_.java.util -import org.apache.spark.annotation.{DeveloperApi, Experimental, Stable} +import org.apache.spark.annotation.{DeveloperApi, Experimental, Stable, Unstable} import org.apache.spark.sql.{Encoder, Row, RuntimeConfig} import org.apache.spark.sql.types.StructType import org.apache.spark.util.SparkClassUtils @@ -93,6 +93,15 @@ abstract class SparkSession extends Serializable with Closeable { */ def udf: UDFRegistration + /** + * Returns a `StreamingQueryManager` that allows managing all the `StreamingQuery`s active on + * `this`. + * + * @since 2.0.0 + */ + @Unstable + def streams: StreamingQueryManager + /** * Start a new session with isolated SQL configurations, temporary tables, registered functions * are isolated, but sharing the underlying `SparkContext` and cached data. diff --git a/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQueryManager.scala b/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQueryManager.scala new file mode 100644 index 0000000000000..88ba9a493d063 --- /dev/null +++ b/sql/api/src/main/scala/org/apache/spark/sql/api/StreamingQueryManager.scala @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.api + +import _root_.java.util.UUID + +import org.apache.spark.annotation.Evolving +import org.apache.spark.sql.streaming.{StreamingQueryException, StreamingQueryListener} + +/** + * A class to manage all the [[StreamingQuery]] active in a `SparkSession`. + * + * @since 2.0.0 + */ +@Evolving +abstract class StreamingQueryManager { + + /** + * Returns a list of active queries associated with this SQLContext + * + * @since 2.0.0 + */ + def active: Array[_ <: StreamingQuery] + + /** + * Returns the query if there is an active query with the given id, or null. + * + * @since 2.1.0 + */ + def get(id: UUID): StreamingQuery + + /** + * Returns the query if there is an active query with the given id, or null. + * + * @since 2.1.0 + */ + def get(id: String): StreamingQuery + + /** + * Wait until any of the queries on the associated SQLContext has terminated since the creation + * of the context, or since `resetTerminated()` was called. If any query was terminated with an + * exception, then the exception will be thrown. + * + * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either + * return immediately (if the query was terminated by `query.stop()`), or throw the exception + * immediately (if the query was terminated with exception). Use `resetTerminated()` to clear + * past terminations and wait for new terminations. + * + * In the case where multiple queries have terminated since `resetTermination()` was called, if + * any query has terminated with exception, then `awaitAnyTermination()` will throw any of the + * exception. For correctly documenting exceptions across multiple queries, users need to stop + * all of them after any of them terminates with exception, and then check the + * `query.exception()` for each query. + * + * @throws org.apache.spark.sql.streaming.StreamingQueryException + * if any query has terminated with an exception + * @since 2.0.0 + */ + @throws[StreamingQueryException] + def awaitAnyTermination(): Unit + + /** + * Wait until any of the queries on the associated SQLContext has terminated since the creation + * of the context, or since `resetTerminated()` was called. Returns whether any query has + * terminated or not (multiple may have terminated). If any query has terminated with an + * exception, then the exception will be thrown. + * + * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either + * return `true` immediately (if the query was terminated by `query.stop()`), or throw the + * exception immediately (if the query was terminated with exception). Use `resetTerminated()` + * to clear past terminations and wait for new terminations. + * + * In the case where multiple queries have terminated since `resetTermination()` was called, if + * any query has terminated with exception, then `awaitAnyTermination()` will throw any of the + * exception. For correctly documenting exceptions across multiple queries, users need to stop + * all of them after any of them terminates with exception, and then check the + * `query.exception()` for each query. + * + * @throws org.apache.spark.sql.streaming.StreamingQueryException + * if any query has terminated with an exception + * @since 2.0.0 + */ + @throws[StreamingQueryException] + def awaitAnyTermination(timeoutMs: Long): Boolean + + /** + * Forget about past terminated queries so that `awaitAnyTermination()` can be used again to + * wait for new terminations. + * + * @since 2.0.0 + */ + def resetTerminated(): Unit + + /** + * Register a [[org.apache.spark.sql.streaming.StreamingQueryListener]] to receive up-calls for + * life cycle events of [[StreamingQuery]]. + * + * @since 2.0.0 + */ + def addListener(listener: StreamingQueryListener): Unit + + /** + * Deregister a [[org.apache.spark.sql.streaming.StreamingQueryListener]]. + * + * @since 2.0.0 + */ + def removeListener(listener: StreamingQueryListener): Unit + + /** + * List all [[org.apache.spark.sql.streaming.StreamingQueryListener]]s attached to this + * [[StreamingQueryManager]]. + * + * @since 3.0.0 + */ + def listListeners(): Array[StreamingQueryListener] +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala index 983cc24718fd2..eeb46fbf145d7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -229,12 +229,7 @@ class SparkSession private( @Unstable def dataSource: DataSourceRegistration = sessionState.dataSourceRegistration - /** - * Returns a `StreamingQueryManager` that allows managing all the - * `StreamingQuery`s active on `this`. - * - * @since 2.0.0 - */ + /** @inheritdoc */ @Unstable def streams: StreamingQueryManager = sessionState.streamingQueryManager diff --git a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala index 9d6fd2e28dea4..42f6d04466b08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Evolving import org.apache.spark.internal.{Logging, MDC} import org.apache.spark.internal.LogKeys.{CLASS_NAME, QUERY_ID, RUN_ID} -import org.apache.spark.sql.{Dataset, SparkSession} +import org.apache.spark.sql.{api, Dataset, SparkSession} import org.apache.spark.sql.catalyst.catalog.CatalogTable import org.apache.spark.sql.catalyst.streaming.{WriteToStream, WriteToStreamStatement} import org.apache.spark.sql.connector.catalog.{Identifier, SupportsWrite, Table, TableCatalog} @@ -47,7 +47,9 @@ import org.apache.spark.util.{Clock, SystemClock, Utils} @Evolving class StreamingQueryManager private[sql] ( sparkSession: SparkSession, - sqlConf: SQLConf) extends Logging { + sqlConf: SQLConf) + extends api.StreamingQueryManager + with Logging { private[sql] val stateStoreCoordinator = StateStoreCoordinatorRef.forDriver(sparkSession.sparkContext.env) @@ -70,7 +72,7 @@ class StreamingQueryManager private[sql] ( * failed. The exception is the exception of the last failed query. */ @GuardedBy("awaitTerminationLock") - private var lastTerminatedQueryException: Option[StreamingQueryException] = null + private var lastTerminatedQueryException: Option[StreamingQueryException] = _ try { sparkSession.sparkContext.conf.get(STREAMING_QUERY_LISTENERS).foreach { classNames => @@ -90,51 +92,20 @@ class StreamingQueryManager private[sql] ( throw QueryExecutionErrors.registeringStreamingQueryListenerError(e) } - /** - * Returns a list of active queries associated with this SQLContext - * - * @since 2.0.0 - */ + /** @inheritdoc */ def active: Array[StreamingQuery] = activeQueriesSharedLock.synchronized { activeQueries.values.toArray } - /** - * Returns the query if there is an active query with the given id, or null. - * - * @since 2.1.0 - */ + /** @inheritdoc */ def get(id: UUID): StreamingQuery = activeQueriesSharedLock.synchronized { activeQueries.get(id).orNull } - /** - * Returns the query if there is an active query with the given id, or null. - * - * @since 2.1.0 - */ + /** @inheritdoc */ def get(id: String): StreamingQuery = get(UUID.fromString(id)) - /** - * Wait until any of the queries on the associated SQLContext has terminated since the - * creation of the context, or since `resetTerminated()` was called. If any query was terminated - * with an exception, then the exception will be thrown. - * - * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either - * return immediately (if the query was terminated by `query.stop()`), - * or throw the exception immediately (if the query was terminated with exception). Use - * `resetTerminated()` to clear past terminations and wait for new terminations. - * - * In the case where multiple queries have terminated since `resetTermination()` was called, - * if any query has terminated with exception, then `awaitAnyTermination()` will - * throw any of the exception. For correctly documenting exceptions across multiple queries, - * users need to stop all of them after any of them terminates with exception, and then check the - * `query.exception()` for each query. - * - * @throws StreamingQueryException if any query has terminated with an exception - * - * @since 2.0.0 - */ + /** @inheritdoc */ @throws[StreamingQueryException] def awaitAnyTermination(): Unit = { awaitTerminationLock.synchronized { @@ -147,27 +118,7 @@ class StreamingQueryManager private[sql] ( } } - /** - * Wait until any of the queries on the associated SQLContext has terminated since the - * creation of the context, or since `resetTerminated()` was called. Returns whether any query - * has terminated or not (multiple may have terminated). If any query has terminated with an - * exception, then the exception will be thrown. - * - * If a query has terminated, then subsequent calls to `awaitAnyTermination()` will either - * return `true` immediately (if the query was terminated by `query.stop()`), - * or throw the exception immediately (if the query was terminated with exception). Use - * `resetTerminated()` to clear past terminations and wait for new terminations. - * - * In the case where multiple queries have terminated since `resetTermination()` was called, - * if any query has terminated with exception, then `awaitAnyTermination()` will - * throw any of the exception. For correctly documenting exceptions across multiple queries, - * users need to stop all of them after any of them terminates with exception, and then check the - * `query.exception()` for each query. - * - * @throws StreamingQueryException if any query has terminated with an exception - * - * @since 2.0.0 - */ + /** @inheritdoc */ @throws[StreamingQueryException] def awaitAnyTermination(timeoutMs: Long): Boolean = { @@ -187,42 +138,24 @@ class StreamingQueryManager private[sql] ( } } - /** - * Forget about past terminated queries so that `awaitAnyTermination()` can be used again to - * wait for new terminations. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def resetTerminated(): Unit = { awaitTerminationLock.synchronized { lastTerminatedQueryException = null } } - /** - * Register a [[StreamingQueryListener]] to receive up-calls for life cycle events of - * [[StreamingQuery]]. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def addListener(listener: StreamingQueryListener): Unit = { listenerBus.addListener(listener) } - /** - * Deregister a [[StreamingQueryListener]]. - * - * @since 2.0.0 - */ + /** @inheritdoc */ def removeListener(listener: StreamingQueryListener): Unit = { listenerBus.removeListener(listener) } - /** - * List all [[StreamingQueryListener]]s attached to this [[StreamingQueryManager]]. - * - * @since 3.0.0 - */ + /** @inheritdoc */ def listListeners(): Array[StreamingQueryListener] = { listenerBus.listeners.asScala.toArray } From 0c1905951f8c31482b0f5ea334c29c13a83cc3c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B8rn=20J=C3=B8rgensen?= Date: Sat, 28 Sep 2024 08:52:13 +0900 Subject: [PATCH 165/189] [SPARK-49820][PYTHON] Change `raise IOError` to `raise OSError` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Change `raise IOError` to `raise OSError` ### Why are the changes needed? > OSError is the builtin error type used for exceptions that relate to the operating system. > > In Python 3.3, a variety of other exceptions, like WindowsError were aliased to OSError. These aliases remain in place for compatibility with older versions of Python, but may be removed in future versions. > > Prefer using OSError directly, as it is more idiomatic and future-proof. > [RUFF rule](https://docs.astral.sh/ruff/rules/os-error-alias/) [Python OSError](https://docs.python.org/3/library/exceptions.html#OSError) ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48287 from bjornjorgensen/IOError-to--OSError. Authored-by: Bjørn Jørgensen Signed-off-by: Hyukjin Kwon --- python/pyspark/install.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyspark/install.py b/python/pyspark/install.py index 90b0150b0a8ca..ba67a157e964d 100644 --- a/python/pyspark/install.py +++ b/python/pyspark/install.py @@ -163,7 +163,7 @@ def install_spark(dest, spark_version, hadoop_version, hive_version): tar.close() if os.path.exists(package_local_path): os.remove(package_local_path) - raise IOError("Unable to download %s." % pretty_pkg_name) + raise OSError("Unable to download %s." % pretty_pkg_name) def get_preferred_mirrors(): From f9a2077fd32faf63796a68cbb3483b486f220b1c Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Sat, 28 Sep 2024 16:21:30 +0900 Subject: [PATCH 166/189] [SPARK-49810][PYTHON] Extract the preparation of `DataFrame.sort` to parent class ### What changes were proposed in this pull request? Extract the preparation of df.sort to parent class ### Why are the changes needed? deduplicate code, the logics in two classes are similar ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #48282 from zhengruifeng/py_sql_sort. Authored-by: Ruifeng Zheng Signed-off-by: Hyukjin Kwon --- python/pyspark/sql/classic/dataframe.py | 52 +++-------------------- python/pyspark/sql/connect/dataframe.py | 53 ++--------------------- python/pyspark/sql/dataframe.py | 56 +++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 96 deletions(-) diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index 0dd66a9d86545..9f9dedbd38207 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -55,6 +55,7 @@ from pyspark.storagelevel import StorageLevel from pyspark.traceback_utils import SCCallSiteSync from pyspark.sql.column import Column +from pyspark.sql.functions import builtin as F from pyspark.sql.classic.column import _to_seq, _to_list, _to_java_column from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter @@ -873,7 +874,8 @@ def sortWithinPartitions( *cols: Union[int, str, Column, List[Union[int, str, Column]]], **kwargs: Any, ) -> ParentDataFrame: - jdf = self._jdf.sortWithinPartitions(self._sort_cols(cols, kwargs)) + _cols = self._preapare_cols_for_sort(F.col, cols, kwargs) + jdf = self._jdf.sortWithinPartitions(self._jseq(_cols, _to_java_column)) return DataFrame(jdf, self.sparkSession) def sort( @@ -881,7 +883,8 @@ def sort( *cols: Union[int, str, Column, List[Union[int, str, Column]]], **kwargs: Any, ) -> ParentDataFrame: - jdf = self._jdf.sort(self._sort_cols(cols, kwargs)) + _cols = self._preapare_cols_for_sort(F.col, cols, kwargs) + jdf = self._jdf.sort(self._jseq(_cols, _to_java_column)) return DataFrame(jdf, self.sparkSession) orderBy = sort @@ -928,51 +931,6 @@ def _jcols_ordinal(self, *cols: "ColumnOrNameOrOrdinal") -> "JavaObject": _cols.append(c) # type: ignore[arg-type] return self._jseq(_cols, _to_java_column) - def _sort_cols( - self, - cols: Sequence[Union[int, str, Column, List[Union[int, str, Column]]]], - kwargs: Dict[str, Any], - ) -> "JavaObject": - """Return a JVM Seq of Columns that describes the sort order""" - if not cols: - raise PySparkValueError( - errorClass="CANNOT_BE_EMPTY", - messageParameters={"item": "column"}, - ) - if len(cols) == 1 and isinstance(cols[0], list): - cols = cols[0] - - jcols = [] - for c in cols: - if isinstance(c, int) and not isinstance(c, bool): - # ordinal is 1-based - if c > 0: - _c = self[c - 1] - # negative ordinal means sort by desc - elif c < 0: - _c = self[-c - 1].desc() - else: - raise PySparkIndexError( - errorClass="ZERO_INDEX", - messageParameters={}, - ) - else: - _c = c # type: ignore[assignment] - jcols.append(_to_java_column(cast("ColumnOrName", _c))) - - ascending = kwargs.get("ascending", True) - if isinstance(ascending, (bool, int)): - if not ascending: - jcols = [jc.desc() for jc in jcols] - elif isinstance(ascending, list): - jcols = [jc if asc else jc.desc() for asc, jc in zip(ascending, jcols)] - else: - raise PySparkTypeError( - errorClass="NOT_BOOL_OR_LIST", - messageParameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, - ) - return self._jseq(jcols) - def describe(self, *cols: Union[str, List[str]]) -> ParentDataFrame: if len(cols) == 1 and isinstance(cols[0], list): cols = cols[0] # type: ignore[assignment] diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 146cfe11bc502..136fe60532df4 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -739,62 +739,16 @@ def limit(self, num: int) -> ParentDataFrame: def tail(self, num: int) -> List[Row]: return DataFrame(plan.Tail(child=self._plan, limit=num), session=self._session).collect() - def _sort_cols( - self, - cols: Sequence[Union[int, str, Column, List[Union[int, str, Column]]]], - kwargs: Dict[str, Any], - ) -> List[Column]: - """Return a JVM Seq of Columns that describes the sort order""" - if cols is None: - raise PySparkValueError( - errorClass="CANNOT_BE_EMPTY", - messageParameters={"item": "cols"}, - ) - - if len(cols) == 1 and isinstance(cols[0], list): - cols = cols[0] - - _cols: List[Column] = [] - for c in cols: - if isinstance(c, int) and not isinstance(c, bool): - # ordinal is 1-based - if c > 0: - _c = self[c - 1] - # negative ordinal means sort by desc - elif c < 0: - _c = self[-c - 1].desc() - else: - raise PySparkIndexError( - errorClass="ZERO_INDEX", - messageParameters={}, - ) - else: - _c = c # type: ignore[assignment] - _cols.append(F._to_col(cast("ColumnOrName", _c))) - - ascending = kwargs.get("ascending", True) - if isinstance(ascending, (bool, int)): - if not ascending: - _cols = [c.desc() for c in _cols] - elif isinstance(ascending, list): - _cols = [c if asc else c.desc() for asc, c in zip(ascending, _cols)] - else: - raise PySparkTypeError( - errorClass="NOT_BOOL_OR_LIST", - messageParameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, - ) - - return [F._sort_col(c) for c in _cols] - def sort( self, *cols: Union[int, str, Column, List[Union[int, str, Column]]], **kwargs: Any, ) -> ParentDataFrame: + _cols = self._preapare_cols_for_sort(F.col, cols, kwargs) res = DataFrame( plan.Sort( self._plan, - columns=self._sort_cols(cols, kwargs), + columns=[F._sort_col(c) for c in _cols], is_global=True, ), session=self._session, @@ -809,10 +763,11 @@ def sortWithinPartitions( *cols: Union[int, str, Column, List[Union[int, str, Column]]], **kwargs: Any, ) -> ParentDataFrame: + _cols = self._preapare_cols_for_sort(F.col, cols, kwargs) res = DataFrame( plan.Sort( self._plan, - columns=self._sort_cols(cols, kwargs), + columns=[F._sort_col(c) for c in _cols], is_global=False, ), session=self._session, diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index 142034583dbd2..5906108163b46 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -2891,6 +2891,62 @@ def sort( """ ... + def _preapare_cols_for_sort( + self, + _to_col: Callable[[str], Column], + cols: Sequence[Union[int, str, Column, List[Union[int, str, Column]]]], + kwargs: Dict[str, Any], + ) -> Sequence[Column]: + from pyspark.errors import PySparkTypeError, PySparkValueError, PySparkIndexError + + if not cols: + raise PySparkValueError( + errorClass="CANNOT_BE_EMPTY", messageParameters={"item": "cols"} + ) + + if len(cols) == 1 and isinstance(cols[0], list): + cols = cols[0] + + _cols: List[Column] = [] + for c in cols: + if isinstance(c, int) and not isinstance(c, bool): + # ordinal is 1-based + if c > 0: + _cols.append(self[c - 1]) + # negative ordinal means sort by desc + elif c < 0: + _cols.append(self[-c - 1].desc()) + else: + raise PySparkIndexError( + errorClass="ZERO_INDEX", + messageParameters={}, + ) + elif isinstance(c, Column): + _cols.append(c) + elif isinstance(c, str): + _cols.append(_to_col(c)) + else: + raise PySparkTypeError( + errorClass="NOT_COLUMN_OR_INT_OR_STR", + messageParameters={ + "arg_name": "col", + "arg_type": type(c).__name__, + }, + ) + + ascending = kwargs.get("ascending", True) + if isinstance(ascending, (bool, int)): + if not ascending: + _cols = [c.desc() for c in _cols] + elif isinstance(ascending, list): + _cols = [c if asc else c.desc() for asc, c in zip(ascending, _cols)] + else: + raise PySparkTypeError( + errorClass="NOT_COLUMN_OR_INT_OR_STR", + messageParameters={"arg_name": "ascending", "arg_type": type(ascending).__name__}, + ) + return _cols + orderBy = sort @dispatch_df_method From 4c12c78801b8de39020981678ec426af8bea00f3 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Sat, 28 Sep 2024 16:25:00 +0900 Subject: [PATCH 167/189] [SPARK-49814][CONNECT] When Spark Connect Client starts, show the `spark version` of the `connect server` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? The pr aims to show the spark version of the connect server when Spark Connect Client starts. ### Why are the changes needed? With the gradual popularize of Spark Connect module, when the Spark Connect client starts, explicitly displaying the spark version of the `connect server`, will reduce confusion for users during execution, such as the new version having some features. However, if it connects to an old version and encounters some problems, it will have to manually troubleshoot. image ### Does this PR introduce _any_ user-facing change? Yes, Connect‘s end-users can intuitively know the `Spark version` on the `server side` when starting the client, reducing confusion. ### How was this patch tested? Manually check. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48283 from panbingkun/SPARK-49814. Authored-by: panbingkun Signed-off-by: Hyukjin Kwon --- .../scala/org/apache/spark/sql/application/ConnectRepl.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/application/ConnectRepl.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/application/ConnectRepl.scala index 63fa2821a6c6a..bff6db25a21f2 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/application/ConnectRepl.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/application/ConnectRepl.scala @@ -50,8 +50,9 @@ object ConnectRepl { /_/ Type in expressions to have them evaluated. +Spark connect server version %s. Spark session available as 'spark'. - """.format(spark_version) + """ def main(args: Array[String]): Unit = doMain(args) @@ -102,7 +103,7 @@ Spark session available as 'spark'. // Please note that we make ammonite generate classes instead of objects. // Classes tend to have superior serialization behavior when using UDFs. val main = new ammonite.Main( - welcomeBanner = Option(splash), + welcomeBanner = Option(splash.format(spark_version, spark.version)), predefCode = predefCode, replCodeWrapper = ExtendedCodeClassWrapper, scriptCodeWrapper = ExtendedCodeClassWrapper, From 550c2071bf8e1e740e595ae9321ae11015d77917 Mon Sep 17 00:00:00 2001 From: panbingkun Date: Sat, 28 Sep 2024 16:23:06 -0700 Subject: [PATCH 168/189] [SPARK-49822][SQL][TESTS] Update postgres docker image to 17.0 ### What changes were proposed in this pull request? This PR aims to update the `postgres` docker image from `16.3` to `17.0`. ### Why are the changes needed? This will help Apache Spark test the latest postgres. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48291 from panbingkun/SPARK-49822. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- .../apache/spark/sql/jdbc/PostgresIntegrationSuite.scala | 6 +++--- .../apache/spark/sql/jdbc/PostgresKrbIntegrationSuite.scala | 6 +++--- .../spark/sql/jdbc/querytest/GeneratedSubquerySuite.scala | 6 +++--- .../spark/sql/jdbc/querytest/PostgreSQLQueryTestSuite.scala | 6 +++--- .../apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala | 6 +++--- .../apache/spark/sql/jdbc/v2/PostgresNamespaceSuite.scala | 6 +++--- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala index 3076b599ef4ef..071b976f044c3 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala @@ -32,9 +32,9 @@ import org.apache.spark.sql.types._ import org.apache.spark.tags.DockerTest /** - * To run this test suite for a specific version (e.g., postgres:16.4-alpine): + * To run this test suite for a specific version (e.g., postgres:17.0-alpine): * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests * "docker-integration-tests/testOnly org.apache.spark.sql.jdbc.PostgresIntegrationSuite" * }}} @@ -42,7 +42,7 @@ import org.apache.spark.tags.DockerTest @DockerTest class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite { override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresKrbIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresKrbIntegrationSuite.scala index 5acb6423bbd9b..62f9c6e0256f3 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresKrbIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresKrbIntegrationSuite.scala @@ -25,9 +25,9 @@ import org.apache.spark.sql.execution.datasources.jdbc.connection.SecureConnecti import org.apache.spark.tags.DockerTest /** - * To run this test suite for a specific version (e.g., postgres:16.4-alpine): + * To run this test suite for a specific version (e.g., postgres:17.0-alpine): * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests * "docker-integration-tests/testOnly *PostgresKrbIntegrationSuite" * }}} @@ -38,7 +38,7 @@ class PostgresKrbIntegrationSuite extends DockerKrbJDBCIntegrationSuite { override protected val keytabFileName = "postgres.keytab" override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/GeneratedSubquerySuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/GeneratedSubquerySuite.scala index 8d367f476403f..a79bbf39a71b8 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/GeneratedSubquerySuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/GeneratedSubquerySuite.scala @@ -28,9 +28,9 @@ import org.apache.spark.tags.DockerTest /** * This suite is used to generate subqueries, and test Spark against Postgres. - * To run this test suite for a specific version (e.g., postgres:16.4-alpine): + * To run this test suite for a specific version (e.g., postgres:17.0-alpine): * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests * "docker-integration-tests/testOnly org.apache.spark.sql.jdbc.GeneratedSubquerySuite" * }}} @@ -39,7 +39,7 @@ import org.apache.spark.tags.DockerTest class GeneratedSubquerySuite extends DockerJDBCIntegrationSuite with QueryGeneratorHelper { override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/PostgreSQLQueryTestSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/PostgreSQLQueryTestSuite.scala index f3a08541365c1..80ba35df6c893 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/PostgreSQLQueryTestSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/querytest/PostgreSQLQueryTestSuite.scala @@ -30,9 +30,9 @@ import org.apache.spark.tags.DockerTest * confidence, and you won't have to manually verify the golden files generated with your test. * 2. Add this line to your .sql file: --ONLY_IF spark * - * Note: To run this test suite for a specific version (e.g., postgres:16.4-alpine): + * Note: To run this test suite for a specific version (e.g., postgres:17.0-alpine): * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests * "testOnly org.apache.spark.sql.jdbc.PostgreSQLQueryTestSuite" * }}} @@ -45,7 +45,7 @@ class PostgreSQLQueryTestSuite extends CrossDbmsQueryTestSuite { protected val customInputFilePath: String = new File(inputFilePath, "subquery").getAbsolutePath override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala index 850391e8dc33c..6bb415a928837 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala @@ -28,9 +28,9 @@ import org.apache.spark.sql.types._ import org.apache.spark.tags.DockerTest /** - * To run this test suite for a specific version (e.g., postgres:16.4-alpine) + * To run this test suite for a specific version (e.g., postgres:17.0-alpine) * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests "testOnly *v2.PostgresIntegrationSuite" * }}} */ @@ -38,7 +38,7 @@ import org.apache.spark.tags.DockerTest class PostgresIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTest { override val catalogName: String = "postgresql" override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresNamespaceSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresNamespaceSuite.scala index 665746f1d5770..6d4f1cc2fd3fc 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresNamespaceSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresNamespaceSuite.scala @@ -26,16 +26,16 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.tags.DockerTest /** - * To run this test suite for a specific version (e.g., postgres:16.4-alpine): + * To run this test suite for a specific version (e.g., postgres:17.0-alpine): * {{{ - * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:16.4-alpine + * ENABLE_DOCKER_INTEGRATION_TESTS=1 POSTGRES_DOCKER_IMAGE_NAME=postgres:17.0-alpine * ./build/sbt -Pdocker-integration-tests "testOnly *v2.PostgresNamespaceSuite" * }}} */ @DockerTest class PostgresNamespaceSuite extends DockerJDBCIntegrationSuite with V2JDBCNamespaceTest { override val db = new DatabaseOnDocker { - override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:16.4-alpine") + override val imageName = sys.env.getOrElse("POSTGRES_DOCKER_IMAGE_NAME", "postgres:17.0-alpine") override val env = Map( "POSTGRES_PASSWORD" -> "rootpass" ) From 47d2c9ca064e9d80a444d21cfac47ca334230242 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Sat, 28 Sep 2024 16:27:13 -0700 Subject: [PATCH 169/189] [SPARK-49712][SQL] Remove encoderFor from connect-client-jvm ### What changes were proposed in this pull request? This PR removes `sql.encoderFor` from the connect-client-jvm module and replaces it by `AgnosticEncoders.agnosticEncoderFor`. ### Why are the changes needed? It will cause a clash when we swap the interface and the implementation. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48266 from hvanhovell/SPARK-49712. Authored-by: Herman van Hovell Signed-off-by: Dongjoon Hyun --- .../main/scala/org/apache/spark/sql/Dataset.scala | 10 +++++----- .../apache/spark/sql/KeyValueGroupedDataset.scala | 14 +++++++------- .../spark/sql/RelationalGroupedDataset.scala | 7 ++++++- .../scala/org/apache/spark/sql/SparkSession.scala | 4 ++-- .../spark/sql/internal/UdfToProtoUtils.scala | 10 +++++----- .../main/scala/org/apache/spark/sql/package.scala | 6 ------ .../apache/spark/sql/SQLImplicitsTestSuite.scala | 3 ++- .../connect/client/arrow/ArrowEncoderSuite.scala | 8 ++++---- 8 files changed, 31 insertions(+), 31 deletions(-) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala index d2877ccaf06c9..6bae04ef80231 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -143,7 +143,7 @@ class Dataset[T] private[sql] ( // Make sure we don't forget to set plan id. assert(plan.getRoot.getCommon.hasPlanId) - private[sql] val agnosticEncoder: AgnosticEncoder[T] = encoderFor(encoder) + private[sql] val agnosticEncoder: AgnosticEncoder[T] = agnosticEncoderFor(encoder) override def toString: String = { try { @@ -437,7 +437,7 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ def select[U1](c1: TypedColumn[T, U1]): Dataset[U1] = { - val encoder = encoderFor(c1.encoder) + val encoder = agnosticEncoderFor(c1.encoder) val col = if (encoder.schema == encoder.dataType) { functions.inline(functions.array(c1)) } else { @@ -452,7 +452,7 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = { - val encoder = ProductEncoder.tuple(columns.map(c => encoderFor(c.encoder))) + val encoder = ProductEncoder.tuple(columns.map(c => agnosticEncoderFor(c.encoder))) selectUntyped(encoder, columns) } @@ -526,7 +526,7 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ def groupByKey[K: Encoder](func: T => K): KeyValueGroupedDataset[K, T] = { - KeyValueGroupedDatasetImpl[K, T](this, encoderFor[K], func) + KeyValueGroupedDatasetImpl[K, T](this, agnosticEncoderFor[K], func) } /** @inheritdoc */ @@ -881,7 +881,7 @@ class Dataset[T] private[sql] ( /** @inheritdoc */ def mapPartitions[U: Encoder](func: Iterator[T] => Iterator[U]): Dataset[U] = { - val outputEncoder = encoderFor[U] + val outputEncoder = agnosticEncoderFor[U] val udf = SparkUserDefinedFunction( function = func, inputEncoders = agnosticEncoder :: Nil, diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala index 6bf2518901470..63b5f27c4745e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala @@ -25,7 +25,7 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.api.java.function._ import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.ProductEncoder +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, ProductEncoder} import org.apache.spark.sql.connect.ConnectConversions._ import org.apache.spark.sql.connect.common.UdfUtils import org.apache.spark.sql.expressions.SparkUserDefinedFunction @@ -398,7 +398,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( new KeyValueGroupedDatasetImpl[L, V, IK, IV]( sparkSession, plan, - encoderFor[L], + agnosticEncoderFor[L], ivEncoder, vEncoder, groupingExprs, @@ -412,7 +412,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( plan, kEncoder, ivEncoder, - encoderFor[W], + agnosticEncoderFor[W], groupingExprs, valueMapFunc .map(_.andThen(valueFunc)) @@ -430,7 +430,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( f: (K, Iterator[V]) => IterableOnce[U]): Dataset[U] = { // Apply mapValues changes to the udf val nf = UDFAdaptors.flatMapGroupsWithMappedValues(f, valueMapFunc) - val outputEncoder = encoderFor[U] + val outputEncoder = agnosticEncoderFor[U] sparkSession.newDataset[U](outputEncoder) { builder => builder.getGroupMapBuilder .setInput(plan.getRoot) @@ -446,7 +446,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( val otherImpl = other.asInstanceOf[KeyValueGroupedDatasetImpl[K, U, _, Any]] // Apply mapValues changes to the udf val nf = UDFAdaptors.coGroupWithMappedValues(f, valueMapFunc, otherImpl.valueMapFunc) - val outputEncoder = encoderFor[R] + val outputEncoder = agnosticEncoderFor[R] sparkSession.newDataset[R](outputEncoder) { builder => builder.getCoGroupMapBuilder .setInput(plan.getRoot) @@ -461,7 +461,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( override protected def aggUntyped(columns: TypedColumn[_, _]*): Dataset[_] = { // TODO(SPARK-43415): For each column, apply the valueMap func first... - val rEnc = ProductEncoder.tuple(kEncoder +: columns.map(c => encoderFor(c.encoder))) + val rEnc = ProductEncoder.tuple(kEncoder +: columns.map(c => agnosticEncoderFor(c.encoder))) sparkSession.newDataset(rEnc) { builder => builder.getAggregateBuilder .setInput(plan.getRoot) @@ -501,7 +501,7 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV]( null } - val outputEncoder = encoderFor[U] + val outputEncoder = agnosticEncoderFor[U] val nf = UDFAdaptors.flatMapGroupsWithStateWithMappedValues(func, valueMapFunc) sparkSession.newDataset[U](outputEncoder) { builder => diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala index 14ceb3f4bb144..5bded40b0d132 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql import scala.jdk.CollectionConverters._ import org.apache.spark.connect.proto +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.agnosticEncoderFor import org.apache.spark.sql.connect.ConnectConversions._ /** @@ -82,7 +83,11 @@ class RelationalGroupedDataset private[sql] ( /** @inheritdoc */ def as[K: Encoder, T: Encoder]: KeyValueGroupedDataset[K, T] = { - KeyValueGroupedDatasetImpl[K, T](df, encoderFor[K], encoderFor[T], groupingExprs) + KeyValueGroupedDatasetImpl[K, T]( + df, + agnosticEncoderFor[K], + agnosticEncoderFor[T], + groupingExprs) } /** @inheritdoc */ diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala index b31670c1da57e..222b5ea79508e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -36,7 +36,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalog.Catalog import org.apache.spark.sql.catalyst.{JavaTypeInference, ScalaReflection} import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, RowEncoder} -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{BoxedLongEncoder, UnboundRowEncoder} +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, BoxedLongEncoder, UnboundRowEncoder} import org.apache.spark.sql.connect.client.{ClassFinder, CloseableIterator, SparkConnectClient, SparkResult} import org.apache.spark.sql.connect.client.SparkConnectClient.Configuration import org.apache.spark.sql.connect.client.arrow.ArrowSerializer @@ -136,7 +136,7 @@ class SparkSession private[sql] ( /** @inheritdoc */ def createDataset[T: Encoder](data: Seq[T]): Dataset[T] = { - createDataset(encoderFor[T], data.iterator) + createDataset(agnosticEncoderFor[T], data.iterator) } /** @inheritdoc */ diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/UdfToProtoUtils.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/UdfToProtoUtils.scala index 85ce2cb820437..409c43f480b8e 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/UdfToProtoUtils.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/internal/UdfToProtoUtils.scala @@ -25,9 +25,9 @@ import com.google.protobuf.ByteString import org.apache.spark.SparkException import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, RowEncoder} +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.agnosticEncoderFor import org.apache.spark.sql.connect.common.DataTypeProtoConverter.toConnectProtoType import org.apache.spark.sql.connect.common.UdfPacket -import org.apache.spark.sql.encoderFor import org.apache.spark.sql.expressions.{SparkUserDefinedFunction, UserDefinedAggregator, UserDefinedFunction} import org.apache.spark.util.{ClosureCleaner, SparkClassUtils, SparkSerDeUtils} @@ -79,12 +79,12 @@ private[sql] object UdfToProtoUtils { udf match { case f: SparkUserDefinedFunction => val outputEncoder = f.outputEncoder - .map(e => encoderFor(e)) + .map(e => agnosticEncoderFor(e)) .getOrElse(RowEncoder.encoderForDataType(f.dataType, lenient = false)) val inputEncoders = if (f.inputEncoders.forall(_.isEmpty)) { Nil // Java UDFs have no bindings for their inputs. } else { - f.inputEncoders.map(e => encoderFor(e.get)) // TODO support Any and UnboundRow. + f.inputEncoders.map(e => agnosticEncoderFor(e.get)) // TODO support Any and UnboundRow. } inputEncoders.foreach(e => protoUdf.addInputTypes(toConnectProtoType(e.dataType))) protoUdf @@ -93,8 +93,8 @@ private[sql] object UdfToProtoUtils { .setAggregate(false) f.givenName.foreach(invokeUdf.setFunctionName) case f: UserDefinedAggregator[_, _, _] => - val outputEncoder = encoderFor(f.aggregator.outputEncoder) - val inputEncoder = encoderFor(f.inputEncoder) + val outputEncoder = agnosticEncoderFor(f.aggregator.outputEncoder) + val inputEncoder = agnosticEncoderFor(f.inputEncoder) protoUdf .setPayload(toUdfPacketBytes(f.aggregator, inputEncoder :: Nil, outputEncoder)) .addInputTypes(toConnectProtoType(inputEncoder.dataType)) diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala index 556b472283a37..ada94b76fcbcd 100644 --- a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala +++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/package.scala @@ -17,12 +17,6 @@ package org.apache.spark -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder - package object sql { type DataFrame = Dataset[Row] - - private[sql] def encoderFor[E: Encoder]: AgnosticEncoder[E] = { - implicitly[Encoder[E]].asInstanceOf[AgnosticEncoder[E]] - } } diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SQLImplicitsTestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SQLImplicitsTestSuite.scala index 57342e12fcb51..b3b8020b1e4c7 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SQLImplicitsTestSuite.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SQLImplicitsTestSuite.scala @@ -26,6 +26,7 @@ import org.apache.arrow.memory.RootAllocator import org.apache.commons.lang3.SystemUtils import org.scalatest.BeforeAndAfterAll +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.agnosticEncoderFor import org.apache.spark.sql.connect.client.SparkConnectClient import org.apache.spark.sql.connect.client.arrow.{ArrowDeserializers, ArrowSerializer} import org.apache.spark.sql.test.ConnectFunSuite @@ -55,7 +56,7 @@ class SQLImplicitsTestSuite extends ConnectFunSuite with BeforeAndAfterAll { import org.apache.spark.util.ArrayImplicits._ import spark.implicits._ def testImplicit[T: Encoder](expected: T): Unit = { - val encoder = encoderFor[T] + val encoder = agnosticEncoderFor[T] val allocator = new RootAllocator() try { val batch = ArrowSerializer.serialize( diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala index 5397dae9dcc5f..7176c582d0bbc 100644 --- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala +++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala @@ -30,11 +30,11 @@ import org.apache.arrow.memory.{BufferAllocator, RootAllocator} import org.apache.arrow.vector.VarBinaryVector import org.scalatest.BeforeAndAfterAll -import org.apache.spark.{sql, SparkRuntimeException, SparkUnsupportedOperationException} +import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException} import org.apache.spark.sql.{AnalysisException, Encoders, Row} import org.apache.spark.sql.catalyst.{DefinedByConstructorParams, JavaTypeInference, ScalaReflection} import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, Codec, OuterScopes} -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{BinaryEncoder, BoxedBooleanEncoder, BoxedByteEncoder, BoxedDoubleEncoder, BoxedFloatEncoder, BoxedIntEncoder, BoxedLongEncoder, BoxedShortEncoder, CalendarIntervalEncoder, DateEncoder, DayTimeIntervalEncoder, EncoderField, InstantEncoder, IterableEncoder, JavaDecimalEncoder, LocalDateEncoder, LocalDateTimeEncoder, NullEncoder, PrimitiveBooleanEncoder, PrimitiveByteEncoder, PrimitiveDoubleEncoder, PrimitiveFloatEncoder, PrimitiveIntEncoder, PrimitiveLongEncoder, PrimitiveShortEncoder, RowEncoder, ScalaDecimalEncoder, StringEncoder, TimestampEncoder, TransformingEncoder, UDTEncoder, YearMonthIntervalEncoder} +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, BinaryEncoder, BoxedBooleanEncoder, BoxedByteEncoder, BoxedDoubleEncoder, BoxedFloatEncoder, BoxedIntEncoder, BoxedLongEncoder, BoxedShortEncoder, CalendarIntervalEncoder, DateEncoder, DayTimeIntervalEncoder, EncoderField, InstantEncoder, IterableEncoder, JavaDecimalEncoder, LocalDateEncoder, LocalDateTimeEncoder, NullEncoder, PrimitiveBooleanEncoder, PrimitiveByteEncoder, PrimitiveDoubleEncoder, PrimitiveFloatEncoder, PrimitiveIntEncoder, PrimitiveLongEncoder, PrimitiveShortEncoder, RowEncoder, ScalaDecimalEncoder, StringEncoder, TimestampEncoder, TransformingEncoder, UDTEncoder, YearMonthIntervalEncoder} import org.apache.spark.sql.catalyst.encoders.RowEncoder.{encoderFor => toRowEncoder} import org.apache.spark.sql.catalyst.util.{DateFormatter, SparkStringUtils, TimestampFormatter} import org.apache.spark.sql.catalyst.util.DateTimeConstants.MICROS_PER_SECOND @@ -770,7 +770,7 @@ class ArrowEncoderSuite extends ConnectFunSuite with BeforeAndAfterAll { } test("java serialization") { - val encoder = sql.encoderFor(Encoders.javaSerialization[(Int, String)]) + val encoder = agnosticEncoderFor(Encoders.javaSerialization[(Int, String)]) roundTripAndCheckIdentical(encoder) { () => Iterator.tabulate(10)(i => (i, "itr_" + i)) } @@ -778,7 +778,7 @@ class ArrowEncoderSuite extends ConnectFunSuite with BeforeAndAfterAll { test("kryo serialization") { val e = intercept[SparkRuntimeException] { - val encoder = sql.encoderFor(Encoders.kryo[(Int, String)]) + val encoder = agnosticEncoderFor(Encoders.kryo[(Int, String)]) roundTripAndCheckIdentical(encoder) { () => Iterator.tabulate(10)(i => (i, "itr_" + i)) } From 8dfecc1463ff0c2a3a18e7a4409736344c2dc3b8 Mon Sep 17 00:00:00 2001 From: Herman van Hovell Date: Sat, 28 Sep 2024 16:30:15 -0700 Subject: [PATCH 170/189] [SPARK-49434][SPARK-49435][CONNECT][SQL] Move aggregators to sql/api ### What changes were proposed in this pull request? This PR moves all user facing Aggregators from sql/core to sql/api. ### Why are the changes needed? We are create a unifies Scala SQL interface. This is part of that effort. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48267 from hvanhovell/SPARK-49434. Authored-by: Herman van Hovell Signed-off-by: Dongjoon Hyun --- project/MimaExcludes.scala | 5 ++++ .../spark/sql/expressions/javalang/typed.java | 10 +++---- .../sql/expressions/ReduceAggregator.scala | 16 +++++------ .../sql/expressions/scalalang/typed.scala | 4 +-- .../sql/internal}/typedaggregators.scala | 27 +++++++++---------- ...ColumnNodeToExpressionConverterSuite.scala | 2 +- 6 files changed, 31 insertions(+), 33 deletions(-) rename sql/{core => api}/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java (88%) rename sql/{core => api}/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala (82%) rename sql/{core => api}/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala (94%) rename sql/{core/src/main/scala/org/apache/spark/sql/execution/aggregate => api/src/main/scala/org/apache/spark/sql/internal}/typedaggregators.scala (81%) diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index 41f547a43b698..2b3d76eb0c2c3 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -184,6 +184,11 @@ object MimaExcludes { ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.avro.functions$"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.protobuf.functions"), ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.protobuf.functions$"), + + // SPARK-49434: Move aggregators to sql/api + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.expressions.javalang.typed"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.expressions.scalalang.typed"), + ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.expressions.scalalang.typed$"), ) ++ loggingExcludes("org.apache.spark.sql.DataFrameReader") ++ loggingExcludes("org.apache.spark.sql.streaming.DataStreamReader") ++ loggingExcludes("org.apache.spark.sql.SparkSession#Builder") diff --git a/sql/core/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java b/sql/api/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java similarity index 88% rename from sql/core/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java rename to sql/api/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java index e1e4ba4c8e0dc..91a1231ec0303 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java +++ b/sql/api/src/main/java/org/apache/spark/sql/expressions/javalang/typed.java @@ -19,13 +19,13 @@ import org.apache.spark.api.java.function.MapFunction; import org.apache.spark.sql.TypedColumn; -import org.apache.spark.sql.execution.aggregate.TypedAverage; -import org.apache.spark.sql.execution.aggregate.TypedCount; -import org.apache.spark.sql.execution.aggregate.TypedSumDouble; -import org.apache.spark.sql.execution.aggregate.TypedSumLong; +import org.apache.spark.sql.internal.TypedAverage; +import org.apache.spark.sql.internal.TypedCount; +import org.apache.spark.sql.internal.TypedSumDouble; +import org.apache.spark.sql.internal.TypedSumLong; /** - * Type-safe functions available for {@link org.apache.spark.sql.Dataset} operations in Java. + * Type-safe functions available for {@link org.apache.spark.sql.api.Dataset} operations in Java. * * Scala users should use {@link org.apache.spark.sql.expressions.scalalang.typed}. * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala b/sql/api/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala similarity index 82% rename from sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala rename to sql/api/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala index 192b5bf65c4c5..9d98d1a98b00d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/expressions/ReduceAggregator.scala @@ -18,19 +18,17 @@ package org.apache.spark.sql.expressions import org.apache.spark.SparkException -import org.apache.spark.sql.Encoder -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder -import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{PrimitiveBooleanEncoder, ProductEncoder} +import org.apache.spark.sql.{Encoder, Encoders} /** * An aggregator that uses a single associative and commutative reduce function. This reduce - * function can be used to go through all input values and reduces them to a single value. - * If there is no input, a null value is returned. + * function can be used to go through all input values and reduces them to a single value. If + * there is no input, a null value is returned. * * This class currently assumes there is at least one input row. */ private[sql] class ReduceAggregator[T: Encoder](func: (T, T) => T) - extends Aggregator[T, (Boolean, T), T] { + extends Aggregator[T, (Boolean, T), T] { @transient private val encoder = implicitly[Encoder[T]] @@ -47,10 +45,8 @@ private[sql] class ReduceAggregator[T: Encoder](func: (T, T) => T) override def zero: (Boolean, T) = (false, _zero.asInstanceOf[T]) - override def bufferEncoder: Encoder[(Boolean, T)] = { - ProductEncoder.tuple(Seq(PrimitiveBooleanEncoder, encoder.asInstanceOf[AgnosticEncoder[T]])) - .asInstanceOf[Encoder[(Boolean, T)]] - } + override def bufferEncoder: Encoder[(Boolean, T)] = + Encoders.tuple(Encoders.scalaBoolean, encoder) override def outputEncoder: Encoder[T] = encoder diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala b/sql/api/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala similarity index 94% rename from sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala rename to sql/api/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala index 8d17edd42442e..9ea3ab8cd4e1c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/expressions/scalalang/typed.scala @@ -17,8 +17,8 @@ package org.apache.spark.sql.expressions.scalalang -import org.apache.spark.sql._ -import org.apache.spark.sql.execution.aggregate._ +import org.apache.spark.sql.TypedColumn +import org.apache.spark.sql.internal.{TypedAverage, TypedCount, TypedSumDouble, TypedSumLong} /** * Type-safe functions available for `Dataset` operations in Scala. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/typedaggregators.scala b/sql/api/src/main/scala/org/apache/spark/sql/internal/typedaggregators.scala similarity index 81% rename from sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/typedaggregators.scala rename to sql/api/src/main/scala/org/apache/spark/sql/internal/typedaggregators.scala index b6550bf3e4aac..aabb3a6f00fd5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/typedaggregators.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/internal/typedaggregators.scala @@ -15,26 +15,24 @@ * limitations under the License. */ -package org.apache.spark.sql.execution.aggregate +package org.apache.spark.sql.internal import org.apache.spark.api.java.function.MapFunction -import org.apache.spark.sql.{Encoder, TypedColumn} -import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.{Encoder, Encoders, TypedColumn} import org.apache.spark.sql.expressions.Aggregator //////////////////////////////////////////////////////////////////////////////////////////////////// // This file defines internal implementations for aggregators. //////////////////////////////////////////////////////////////////////////////////////////////////// - class TypedSumDouble[IN](val f: IN => Double) extends Aggregator[IN, Double, Double] { override def zero: Double = 0.0 override def reduce(b: Double, a: IN): Double = b + f(a) override def merge(b1: Double, b2: Double): Double = b1 + b2 override def finish(reduction: Double): Double = reduction - override def bufferEncoder: Encoder[Double] = ExpressionEncoder[Double]() - override def outputEncoder: Encoder[Double] = ExpressionEncoder[Double]() + override def bufferEncoder: Encoder[Double] = Encoders.scalaDouble + override def outputEncoder: Encoder[Double] = Encoders.scalaDouble // Java api support def this(f: MapFunction[IN, java.lang.Double]) = this((x: IN) => f.call(x).asInstanceOf[Double]) @@ -44,15 +42,14 @@ class TypedSumDouble[IN](val f: IN => Double) extends Aggregator[IN, Double, Dou } } - class TypedSumLong[IN](val f: IN => Long) extends Aggregator[IN, Long, Long] { override def zero: Long = 0L override def reduce(b: Long, a: IN): Long = b + f(a) override def merge(b1: Long, b2: Long): Long = b1 + b2 override def finish(reduction: Long): Long = reduction - override def bufferEncoder: Encoder[Long] = ExpressionEncoder[Long]() - override def outputEncoder: Encoder[Long] = ExpressionEncoder[Long]() + override def bufferEncoder: Encoder[Long] = Encoders.scalaLong + override def outputEncoder: Encoder[Long] = Encoders.scalaLong // Java api support def this(f: MapFunction[IN, java.lang.Long]) = this((x: IN) => f.call(x).asInstanceOf[Long]) @@ -62,7 +59,6 @@ class TypedSumLong[IN](val f: IN => Long) extends Aggregator[IN, Long, Long] { } } - class TypedCount[IN](val f: IN => Any) extends Aggregator[IN, Long, Long] { override def zero: Long = 0 override def reduce(b: Long, a: IN): Long = { @@ -71,8 +67,8 @@ class TypedCount[IN](val f: IN => Any) extends Aggregator[IN, Long, Long] { override def merge(b1: Long, b2: Long): Long = b1 + b2 override def finish(reduction: Long): Long = reduction - override def bufferEncoder: Encoder[Long] = ExpressionEncoder[Long]() - override def outputEncoder: Encoder[Long] = ExpressionEncoder[Long]() + override def bufferEncoder: Encoder[Long] = Encoders.scalaLong + override def outputEncoder: Encoder[Long] = Encoders.scalaLong // Java api support def this(f: MapFunction[IN, Object]) = this((x: IN) => f.call(x).asInstanceOf[Any]) @@ -81,7 +77,6 @@ class TypedCount[IN](val f: IN => Any) extends Aggregator[IN, Long, Long] { } } - class TypedAverage[IN](val f: IN => Double) extends Aggregator[IN, (Double, Long), Double] { override def zero: (Double, Long) = (0.0, 0L) override def reduce(b: (Double, Long), a: IN): (Double, Long) = (f(a) + b._1, 1 + b._2) @@ -90,8 +85,10 @@ class TypedAverage[IN](val f: IN => Double) extends Aggregator[IN, (Double, Long (b1._1 + b2._1, b1._2 + b2._2) } - override def bufferEncoder: Encoder[(Double, Long)] = ExpressionEncoder[(Double, Long)]() - override def outputEncoder: Encoder[Double] = ExpressionEncoder[Double]() + override def bufferEncoder: Encoder[(Double, Long)] = + Encoders.tuple(Encoders.scalaDouble, Encoders.scalaLong) + + override def outputEncoder: Encoder[Double] = Encoders.scalaDouble // Java api support def this(f: MapFunction[IN, java.lang.Double]) = this((x: IN) => f.call(x).asInstanceOf[Double]) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/internal/ColumnNodeToExpressionConverterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/internal/ColumnNodeToExpressionConverterSuite.scala index c993aa8e52031..76fcdfc380950 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/internal/ColumnNodeToExpressionConverterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/internal/ColumnNodeToExpressionConverterSuite.scala @@ -324,7 +324,7 @@ class ColumnNodeToExpressionConverterSuite extends SparkFunSuite { a.asInstanceOf[AgnosticEncoder[Any]] test("udf") { - val int2LongSum = new aggregate.TypedSumLong[Int]((i: Int) => i.toLong) + val int2LongSum = new TypedSumLong[Int]((i: Int) => i.toLong) val bufferEncoder = encoderFor(int2LongSum.bufferEncoder) val outputEncoder = encoderFor(int2LongSum.outputEncoder) val bufferAttrs = bufferEncoder.namedExpressions.map { From 039fd13eacb1cef835045e3a60cebf958589e1a2 Mon Sep 17 00:00:00 2001 From: Yuming Wang Date: Sat, 28 Sep 2024 19:45:52 -0700 Subject: [PATCH 171/189] [SPARK-49749][CORE] Change log level to debug in BlockManagerInfo ### What changes were proposed in this pull request? This PR changes the log level to debug in `BlockManagerInfo`. ### Why are the changes needed? Before this PR: Logging in `BlockManagerMasterEndpoint` uses 3.25% of the CPU and generates 60.5% of the logs. image ``` cat spark.20240921-09.log | grep "in memory on" | wc -l 8587851 cat spark.20240921-09.log | wc -l 14185544 ``` After this PR: image ``` cat spark.20240926-09.log | grep "in memory on" | wc -l 0 cat spark.20240926-09.log | wc -l 2224037 ``` ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48197 from wangyum/SPARK-49749. Authored-by: Yuming Wang Signed-off-by: Dongjoon Hyun --- .../spark/storage/BlockManagerMasterEndpoint.scala | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala index 73f89ea0e86e5..fc4e6e771aad7 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala @@ -1059,13 +1059,13 @@ private[spark] class BlockManagerInfo( _blocks.put(blockId, blockStatus) _remainingMem -= memSize if (blockExists) { - logInfo(log"Updated ${MDC(BLOCK_ID, blockId)} in memory on " + + logDebug(log"Updated ${MDC(BLOCK_ID, blockId)} in memory on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} (current size: " + log"${MDC(CURRENT_MEMORY_SIZE, Utils.bytesToString(memSize))}, original " + log"size: ${MDC(ORIGINAL_MEMORY_SIZE, Utils.bytesToString(originalMemSize))}, " + log"free: ${MDC(FREE_MEMORY_SIZE, Utils.bytesToString(_remainingMem))})") } else { - logInfo(log"Added ${MDC(BLOCK_ID, blockId)} in memory on " + + logDebug(log"Added ${MDC(BLOCK_ID, blockId)} in memory on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} " + log"(size: ${MDC(CURRENT_MEMORY_SIZE, Utils.bytesToString(memSize))}, " + log"free: ${MDC(FREE_MEMORY_SIZE, Utils.bytesToString(_remainingMem))})") @@ -1075,12 +1075,12 @@ private[spark] class BlockManagerInfo( blockStatus = BlockStatus(storageLevel, memSize = 0, diskSize = diskSize) _blocks.put(blockId, blockStatus) if (blockExists) { - logInfo(log"Updated ${MDC(BLOCK_ID, blockId)} on disk on " + + logDebug(log"Updated ${MDC(BLOCK_ID, blockId)} on disk on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} " + log"(current size: ${MDC(CURRENT_DISK_SIZE, Utils.bytesToString(diskSize))}," + log" original size: ${MDC(ORIGINAL_DISK_SIZE, Utils.bytesToString(originalDiskSize))})") } else { - logInfo(log"Added ${MDC(BLOCK_ID, blockId)} on disk on " + + logDebug(log"Added ${MDC(BLOCK_ID, blockId)} on disk on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} (size: " + log"${MDC(CURRENT_DISK_SIZE, Utils.bytesToString(diskSize))})") } @@ -1098,13 +1098,13 @@ private[spark] class BlockManagerInfo( blockStatus.remove(blockId) } if (originalLevel.useMemory) { - logInfo(log"Removed ${MDC(BLOCK_ID, blockId)} on " + + logDebug(log"Removed ${MDC(BLOCK_ID, blockId)} on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} in memory " + log"(size: ${MDC(ORIGINAL_MEMORY_SIZE, Utils.bytesToString(originalMemSize))}, " + log"free: ${MDC(FREE_MEMORY_SIZE, Utils.bytesToString(_remainingMem))})") } if (originalLevel.useDisk) { - logInfo(log"Removed ${MDC(BLOCK_ID, blockId)} on " + + logDebug(log"Removed ${MDC(BLOCK_ID, blockId)} on " + log"${MDC(HOST_PORT, blockManagerId.hostPort)} on disk" + log" (size: ${MDC(ORIGINAL_DISK_SIZE, Utils.bytesToString(originalDiskSize))})") } From 885c3fac724611ca59add984eb0629d32644b56f Mon Sep 17 00:00:00 2001 From: Anish Shrigondekar Date: Mon, 30 Sep 2024 15:02:40 +0900 Subject: [PATCH 172/189] [SPARK-49823][SS] Avoid flush during shutdown in rocksdb close path ### What changes were proposed in this pull request? Avoid flush during shutdown in rocksdb close path ### Why are the changes needed? Without this change, we see sometimes that `cancelAllBackgroundWork` gets hung if there are memtables that need to be flushed. We also don't need to flush in this path, because we only assume that sync flush is required in the commit path. ``` at app//org.rocksdb.RocksDB.cancelAllBackgroundWork(Native Method) at app//org.rocksdb.RocksDB.cancelAllBackgroundWork(RocksDB.java:4053) at app//org.apache.spark.sql.execution.streaming.state.RocksDB.closeDB(RocksDB.scala:1406) at app//org.apache.spark.sql.execution.streaming.state.RocksDB.load(RocksDB.scala:383) ``` ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Verified the config is passed manually in the logs and existing unit tests. Before: ``` sql/core/target/unit-tests.log:141:18:20:06.223 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 0 sql/core/target/unit-tests.log:776:18:20:06.871 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 0 sql/core/target/unit-tests.log:1096:18:20:07.129 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 0 ``` After: ``` sql/core/target/unit-tests.log:6561:18:17:42.723 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 1 sql/core/target/unit-tests.log:6947:18:17:43.035 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 1 sql/core/target/unit-tests.log:7344:18:17:43.313 pool-1-thread-1-ScalaTest-running-RocksDBSuite INFO RocksDB [Thread-17]: [NativeRocksDB-1] Options.avoid_flush_during_shutdown: 1 ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48292 from anishshri-db/task/SPARK-49823. Authored-by: Anish Shrigondekar Signed-off-by: Jungtaek Lim --- .../org/apache/spark/sql/execution/streaming/state/RocksDB.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala index f8d0c8722c3f5..c7f8434e5345b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala @@ -134,6 +134,7 @@ class RocksDB( rocksDbOptions.setTableFormatConfig(tableFormatConfig) rocksDbOptions.setMaxOpenFiles(conf.maxOpenFiles) rocksDbOptions.setAllowFAllocate(conf.allowFAllocate) + rocksDbOptions.setAvoidFlushDuringShutdown(true) rocksDbOptions.setMergeOperator(new StringAppendOperator()) if (conf.boundedMemoryUsage) { From d85e7bc0beb49dd1d894d487cf6a5a02075280dd Mon Sep 17 00:00:00 2001 From: Jovan Pavlovic Date: Mon, 30 Sep 2024 17:41:42 +0800 Subject: [PATCH 173/189] [SPARK-49811][SQL] Rename StringTypeAnyCollation ### What changes were proposed in this pull request? Rename StringTypeAnyCollation to StringTypeWithCaseAccentSensitivity. Name StringTypeAnyCollation is unfortunate, with adding new type of collations it requires ren ### Why are the changes needed? Name StringTypeAnyCollation is unfortunate, with adding new specifier (for example trim specifier) it requires always renaming it to (something like AllCollationExeptTrimCollation) until new collation is implemented in all functions. It gets even more confusing if multiple collations are not supported for some functions. Instead of this naming convention should be only specifiers that are supported and avoid using all. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Just renaming all tests passing. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48265 from jovanpavl-db/rename-string-type-collations. Authored-by: Jovan Pavlovic Signed-off-by: Wenchen Fan --- .../internal/types/AbstractStringType.scala | 7 +- .../sql/catalyst/analysis/TypeCoercion.scala | 5 +- .../expressions/CallMethodViaReflection.scala | 9 +- .../catalyst/expressions/CollationKey.scala | 4 +- .../sql/catalyst/expressions/ExprUtils.scala | 5 +- .../aggregate/datasketchesAggregates.scala | 6 +- .../expressions/collationExpressions.scala | 6 +- .../expressions/collectionOperations.scala | 13 ++- .../catalyst/expressions/csvExpressions.scala | 4 +- .../expressions/datetimeExpressions.scala | 41 ++++--- .../expressions/jsonExpressions.scala | 14 ++- .../expressions/maskExpressions.scala | 10 +- .../expressions/mathExpressions.scala | 8 +- .../spark/sql/catalyst/expressions/misc.scala | 13 ++- .../expressions/numberFormatExpressions.scala | 7 +- .../expressions/regexpExpressions.scala | 18 +-- .../expressions/stringExpressions.scala | 103 ++++++++++-------- .../catalyst/expressions/urlExpressions.scala | 13 ++- .../variant/variantExpressions.scala | 7 +- .../sql/catalyst/expressions/xml/xpath.scala | 6 +- .../catalyst/expressions/xmlExpressions.scala | 4 +- .../analysis/AnsiTypeCoercionSuite.scala | 20 ++-- .../expressions/StringExpressionsSuite.scala | 4 +- .../sql/CollationExpressionWalkerSuite.scala | 51 +++++---- 24 files changed, 218 insertions(+), 160 deletions(-) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala b/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala index dc4ee013fd189..6feb662632763 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/internal/types/AbstractStringType.scala @@ -21,7 +21,7 @@ import org.apache.spark.sql.internal.SqlApiConf import org.apache.spark.sql.types.{AbstractDataType, DataType, StringType} /** - * StringTypeCollated is an abstract class for StringType with collation support. + * AbstractStringType is an abstract class for StringType with collation support. */ abstract class AbstractStringType extends AbstractDataType { override private[sql] def defaultConcreteType: DataType = SqlApiConf.get.defaultStringType @@ -46,9 +46,10 @@ case object StringTypeBinaryLcase extends AbstractStringType { } /** - * Use StringTypeAnyCollation for expressions supporting all possible collation types. + * Use StringTypeWithCaseAccentSensitivity for expressions supporting all collation types (binary + * and ICU) but limited to using case and accent sensitivity specifiers. */ -case object StringTypeAnyCollation extends AbstractStringType { +case object StringTypeWithCaseAccentSensitivity extends AbstractStringType { override private[sql] def acceptsType(other: DataType): Boolean = other.isInstanceOf[StringType] } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala index 5983346ff1e27..e0298b19931c7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala @@ -32,7 +32,8 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.connector.catalog.procedures.BoundProcedure import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, AbstractMapType, AbstractStringType, StringTypeAnyCollation} +import org.apache.spark.sql.internal.types.{AbstractArrayType, AbstractMapType, AbstractStringType, + StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types._ import org.apache.spark.sql.types.UpCastRule.numericPrecedence @@ -438,7 +439,7 @@ abstract class TypeCoercionBase { } case aj @ ArrayJoin(arr, d, nr) - if !AbstractArrayType(StringTypeAnyCollation).acceptsType(arr.dataType) && + if !AbstractArrayType(StringTypeWithCaseAccentSensitivity).acceptsType(arr.dataType) && ArrayType.acceptsType(arr.dataType) => val containsNull = arr.dataType.asInstanceOf[ArrayType].containsNull implicitCast(arr, ArrayType(StringType, containsNull)) match { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala index 13ea8c77c41b4..6aa11b6fd16df 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.ArrayImplicits._ @@ -84,7 +84,7 @@ case class CallMethodViaReflection( errorSubClass = "NON_FOLDABLE_INPUT", messageParameters = Map( "inputName" -> toSQLId("class"), - "inputType" -> toSQLType(StringTypeAnyCollation), + "inputType" -> toSQLType(StringTypeWithCaseAccentSensitivity), "inputExpr" -> toSQLExpr(children.head) ) ) @@ -97,7 +97,7 @@ case class CallMethodViaReflection( errorSubClass = "NON_FOLDABLE_INPUT", messageParameters = Map( "inputName" -> toSQLId("method"), - "inputType" -> toSQLType(StringTypeAnyCollation), + "inputType" -> toSQLType(StringTypeWithCaseAccentSensitivity), "inputExpr" -> toSQLExpr(children(1)) ) ) @@ -114,7 +114,8 @@ case class CallMethodViaReflection( "paramIndex" -> ordinalNumber(idx), "requiredType" -> toSQLType( TypeCollection(BooleanType, ByteType, ShortType, - IntegerType, LongType, FloatType, DoubleType, StringTypeAnyCollation)), + IntegerType, LongType, FloatType, DoubleType, + StringTypeWithCaseAccentSensitivity)), "inputSql" -> toSQLExpr(e), "inputType" -> toSQLType(e.dataType)) ) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CollationKey.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CollationKey.scala index 6e400d026e0ee..28ec8482e5cdd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CollationKey.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CollationKey.scala @@ -19,12 +19,12 @@ package org.apache.spark.sql.catalyst.expressions import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} import org.apache.spark.sql.catalyst.util.CollationFactory -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String case class CollationKey(expr: Expression) extends UnaryExpression with ExpectsInputTypes { - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def dataType: DataType = BinaryType final lazy val collationId: Int = expr.dataType match { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala index 749152f135e92..08cb03edb78b6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.logical.Aggregate import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, CharVarcharUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors} -import org.apache.spark.sql.internal.types.{AbstractMapType, StringTypeAnyCollation} +import org.apache.spark.sql.internal.types.{AbstractMapType, StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types.{DataType, MapType, StringType, StructType, VariantType} import org.apache.spark.unsafe.types.UTF8String @@ -61,7 +61,8 @@ object ExprUtils extends QueryErrorsBase { def convertToMapData(exp: Expression): Map[String, String] = exp match { case m: CreateMap - if AbstractMapType(StringTypeAnyCollation, StringTypeAnyCollation).acceptsType(m.dataType) => + if AbstractMapType(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) + .acceptsType(m.dataType) => val arrayMap = m.eval().asInstanceOf[ArrayBasedMapData] ArrayBasedMapData.toScalaMap(arrayMap).map { case (key, value) => key.toString -> value.toString diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala index 2102428131f64..78bd02d5703cd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes, Expression, import org.apache.spark.sql.catalyst.trees.BinaryLike import org.apache.spark.sql.catalyst.util.CollationFactory import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types.{AbstractDataType, BinaryType, BooleanType, DataType, IntegerType, LongType, StringType, TypeCollection} import org.apache.spark.unsafe.types.UTF8String @@ -105,7 +105,9 @@ case class HllSketchAgg( override def prettyName: String = "hll_sketch_agg" override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(IntegerType, LongType, StringTypeAnyCollation, BinaryType), IntegerType) + Seq( + TypeCollection(IntegerType, LongType, StringTypeWithCaseAccentSensitivity, BinaryType), + IntegerType) override def dataType: DataType = BinaryType diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala index d45ca533f9392..0cff70436db7d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.util.CollationFactory import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ // scalastyle:off line.contains.tab @@ -73,7 +73,7 @@ case class Collate(child: Expression, collationName: String) extends UnaryExpression with ExpectsInputTypes { private val collationId = CollationFactory.collationNameToId(collationName) override def dataType: DataType = StringType(collationId) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override protected def withNewChildInternal( newChild: Expression): Expression = copy(newChild) @@ -111,5 +111,5 @@ case class Collation(child: Expression) val collationName = CollationFactory.fetchCollation(collationId).collationName Literal.create(collationName, SQLConf.get.defaultStringType) } - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala index 5cdd3c7eb62d1..c091d51fc177f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.catalyst.util.DateTimeUtils._ import org.apache.spark.sql.errors.{QueryErrorsBase, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeAnyCollation} +import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types._ import org.apache.spark.sql.util.SQLOpenHashSet import org.apache.spark.unsafe.UTF8StringBuilder @@ -1348,7 +1348,7 @@ case class Reverse(child: Expression) // Input types are utilized by type coercion in ImplicitTypeCasts. override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, ArrayType)) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, ArrayType)) override def dataType: DataType = child.dataType @@ -2134,9 +2134,12 @@ case class ArrayJoin( this(array, delimiter, Some(nullReplacement)) override def inputTypes: Seq[AbstractDataType] = if (nullReplacement.isDefined) { - Seq(AbstractArrayType(StringTypeAnyCollation), StringTypeAnyCollation, StringTypeAnyCollation) + Seq(AbstractArrayType(StringTypeWithCaseAccentSensitivity), + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity) } else { - Seq(AbstractArrayType(StringTypeAnyCollation), StringTypeAnyCollation) + Seq(AbstractArrayType(StringTypeWithCaseAccentSensitivity), + StringTypeWithCaseAccentSensitivity) } override def children: Seq[Expression] = if (nullReplacement.isDefined) { @@ -2857,7 +2860,7 @@ case class Concat(children: Seq[Expression]) extends ComplexTypeMergingExpressio with QueryErrorsBase { private def allowedTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, BinaryType, ArrayType) + Seq(StringTypeWithCaseAccentSensitivity, BinaryType, ArrayType) final override val nodePatterns: Seq[TreePattern] = Seq(CONCAT) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala index cb10440c48328..2f4462c0664f8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.catalyst.util.TypeUtils._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -147,7 +147,7 @@ case class CsvToStructs( converter(parser.parse(csv)) } - override def inputTypes: Seq[AbstractDataType] = StringTypeAnyCollation :: Nil + override def inputTypes: Seq[AbstractDataType] = StringTypeWithCaseAccentSensitivity :: Nil override def prettyName: String = "from_csv" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala index 36bd53001594e..b166d235557fc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils._ import org.apache.spark.sql.catalyst.util.LegacyDateFormats.SIMPLE_DATE_FORMAT import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.sql.types.DayTimeIntervalType.DAY import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} @@ -961,7 +961,8 @@ case class DateFormatClass(left: Expression, right: Expression, timeZoneId: Opti override def dataType: DataType = SQLConf.get.defaultStringType - override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(TimestampType, StringTypeWithCaseAccentSensitivity) override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = copy(timeZoneId = Option(timeZoneId)) @@ -1269,8 +1270,10 @@ abstract class ToTimestamp override def forTimestampNTZ: Boolean = left.dataType == TimestampNTZType override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, DateType, TimestampType, TimestampNTZType), - StringTypeAnyCollation) + Seq(TypeCollection( + StringTypeWithCaseAccentSensitivity, DateType, TimestampType, TimestampNTZType + ), + StringTypeWithCaseAccentSensitivity) override def dataType: DataType = LongType override def nullable: Boolean = if (failOnError) children.exists(_.nullable) else true @@ -1441,7 +1444,8 @@ case class FromUnixTime(sec: Expression, format: Expression, timeZoneId: Option[ override def dataType: DataType = SQLConf.get.defaultStringType override def nullable: Boolean = true - override def inputTypes: Seq[AbstractDataType] = Seq(LongType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(LongType, StringTypeWithCaseAccentSensitivity) override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = copy(timeZoneId = Option(timeZoneId)) @@ -1549,7 +1553,8 @@ case class NextDay( def this(left: Expression, right: Expression) = this(left, right, SQLConf.get.ansiEnabled) - override def inputTypes: Seq[AbstractDataType] = Seq(DateType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(DateType, StringTypeWithCaseAccentSensitivity) override def dataType: DataType = DateType override def nullable: Boolean = true @@ -1760,7 +1765,8 @@ sealed trait UTCTimestamp extends BinaryExpression with ImplicitCastInputTypes w val func: (Long, String) => Long val funcName: String - override def inputTypes: Seq[AbstractDataType] = Seq(TimestampType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(TimestampType, StringTypeWithCaseAccentSensitivity) override def dataType: DataType = TimestampType override def nullSafeEval(time: Any, timezone: Any): Any = { @@ -2100,8 +2106,9 @@ case class ParseToDate( override def inputTypes: Seq[AbstractDataType] = { // Note: ideally this function should only take string input, but we allow more types here to // be backward compatible. - TypeCollection(StringTypeAnyCollation, DateType, TimestampType, TimestampNTZType) +: - format.map(_ => StringTypeAnyCollation).toSeq + TypeCollection( + StringTypeWithCaseAccentSensitivity, DateType, TimestampType, TimestampNTZType) +: + format.map(_ => StringTypeWithCaseAccentSensitivity).toSeq } override protected def withNewChildrenInternal( @@ -2172,10 +2179,10 @@ case class ParseToTimestamp( override def inputTypes: Seq[AbstractDataType] = { // Note: ideally this function should only take string input, but we allow more types here to // be backward compatible. - val types = Seq(StringTypeAnyCollation, DateType, TimestampType, TimestampNTZType) + val types = Seq(StringTypeWithCaseAccentSensitivity, DateType, TimestampType, TimestampNTZType) TypeCollection( (if (dataType.isInstanceOf[TimestampType]) types :+ NumericType else types): _* - ) +: format.map(_ => StringTypeAnyCollation).toSeq + ) +: format.map(_ => StringTypeWithCaseAccentSensitivity).toSeq } override protected def withNewChildrenInternal( @@ -2305,7 +2312,8 @@ case class TruncDate(date: Expression, format: Expression) override def left: Expression = date override def right: Expression = format - override def inputTypes: Seq[AbstractDataType] = Seq(DateType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(DateType, StringTypeWithCaseAccentSensitivity) override def dataType: DataType = DateType override def prettyName: String = "trunc" override val instant = date @@ -2374,7 +2382,8 @@ case class TruncTimestamp( override def left: Expression = format override def right: Expression = timestamp - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation, TimestampType) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCaseAccentSensitivity, TimestampType) override def dataType: TimestampType = TimestampType override def prettyName: String = "date_trunc" override val instant = timestamp @@ -2675,7 +2684,7 @@ case class MakeTimestamp( // casted into decimal safely, we use DecimalType(16, 6) which is wider than DecimalType(10, 0). override def inputTypes: Seq[AbstractDataType] = Seq(IntegerType, IntegerType, IntegerType, IntegerType, IntegerType, DecimalType(16, 6)) ++ - timezone.map(_ => StringTypeAnyCollation) + timezone.map(_ => StringTypeWithCaseAccentSensitivity) override def nullable: Boolean = if (failOnError) children.exists(_.nullable) else true override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = @@ -3122,8 +3131,8 @@ case class ConvertTimezone( override def second: Expression = targetTz override def third: Expression = sourceTs - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation, - StringTypeAnyCollation, TimestampNTZType) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity, TimestampNTZType) override def dataType: DataType = TimestampNTZType override def nullSafeEval(srcTz: Any, tgtTz: Any, micros: Any): Any = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala index 2037eb22fede6..bdcf3f0c1eeab 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{JSON_TO_STRUCT, TreePatt import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.{UTF8String, VariantVal} import org.apache.spark.util.Utils @@ -134,7 +134,7 @@ case class GetJsonObject(json: Expression, path: Expression) override def left: Expression = json override def right: Expression = path override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override def dataType: DataType = SQLConf.get.defaultStringType override def nullable: Boolean = true override def prettyName: String = "get_json_object" @@ -489,7 +489,9 @@ case class JsonTuple(children: Seq[Expression]) throw QueryCompilationErrors.wrongNumArgsError( toSQLId(prettyName), Seq("> 1"), children.length ) - } else if (children.forall(child => StringTypeAnyCollation.acceptsType(child.dataType))) { + } else if ( + children.forall( + child => StringTypeWithCaseAccentSensitivity.acceptsType(child.dataType))) { TypeCheckResult.TypeCheckSuccess } else { DataTypeMismatch( @@ -726,7 +728,7 @@ case class JsonToStructs( converter(parser.parse(json.asInstanceOf[UTF8String])) } - override def inputTypes: Seq[AbstractDataType] = StringTypeAnyCollation :: Nil + override def inputTypes: Seq[AbstractDataType] = StringTypeWithCaseAccentSensitivity :: Nil override def sql: String = schema match { case _: MapType => "entries" @@ -968,7 +970,7 @@ case class SchemaOfJson( case class LengthOfJsonArray(child: Expression) extends UnaryExpression with CodegenFallback with ExpectsInputTypes { - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def dataType: DataType = IntegerType override def nullable: Boolean = true override def prettyName: String = "json_array_length" @@ -1041,7 +1043,7 @@ case class LengthOfJsonArray(child: Expression) extends UnaryExpression case class JsonObjectKeys(child: Expression) extends UnaryExpression with CodegenFallback with ExpectsInputTypes { - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def dataType: DataType = ArrayType(SQLConf.get.defaultStringType) override def nullable: Boolean = true override def prettyName: String = "json_object_keys" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala index c11357352c79a..cb62fa2cc3bd5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.plans.logical.{FunctionSignature, InputParameter} import org.apache.spark.sql.errors.QueryErrorsBase import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types.{AbstractDataType, DataType} import org.apache.spark.unsafe.types.UTF8String @@ -192,8 +192,12 @@ case class Mask( * NumericType, IntegralType, FractionalType. */ override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation, - StringTypeAnyCollation, StringTypeAnyCollation) + Seq( + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity) override def nullable: Boolean = true diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala index ddba820414ae4..e46acf467db22 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.util.{MathUtils, NumberConverter, TypeUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -453,7 +453,7 @@ case class Conv( override def second: Expression = fromBaseExpr override def third: Expression = toBaseExpr override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, IntegerType, IntegerType) + Seq(StringTypeWithCaseAccentSensitivity, IntegerType, IntegerType) override def dataType: DataType = first.dataType override def nullable: Boolean = true @@ -1114,7 +1114,7 @@ case class Hex(child: Expression) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(LongType, BinaryType, StringTypeAnyCollation)) + Seq(TypeCollection(LongType, BinaryType, StringTypeWithCaseAccentSensitivity)) override def dataType: DataType = child.dataType match { case st: StringType => st @@ -1158,7 +1158,7 @@ case class Unhex(child: Expression, failOnError: Boolean = false) def this(expr: Expression) = this(expr, false) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nullable: Boolean = true override def dataType: DataType = BinaryType diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala index 6629f724c4dda..cb846f606632b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.util.{MapData, RandomUUIDGenerator} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.errors.QueryExecutionErrors.raiseError import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -85,7 +85,7 @@ case class RaiseError(errorClass: Expression, errorParms: Expression, dataType: override def foldable: Boolean = false override def nullable: Boolean = true override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, MapType(StringType, StringType)) + Seq(StringTypeWithCaseAccentSensitivity, MapType(StringType, StringType)) override def left: Expression = errorClass override def right: Expression = errorParms @@ -415,7 +415,9 @@ case class AesEncrypt( override def prettyName: String = "aes_encrypt" override def inputTypes: Seq[AbstractDataType] = - Seq(BinaryType, BinaryType, StringTypeAnyCollation, StringTypeAnyCollation, + Seq(BinaryType, BinaryType, + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, BinaryType, BinaryType) override def children: Seq[Expression] = Seq(input, key, mode, padding, iv, aad) @@ -489,7 +491,10 @@ case class AesDecrypt( this(input, key, Literal("GCM")) override def inputTypes: Seq[AbstractDataType] = { - Seq(BinaryType, BinaryType, StringTypeAnyCollation, StringTypeAnyCollation, BinaryType) + Seq(BinaryType, + BinaryType, + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, BinaryType) } override def prettyName: String = "aes_decrypt" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/numberFormatExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/numberFormatExpressions.scala index e914190c06456..5bd2ab6035e10 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/numberFormatExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/numberFormatExpressions.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.util.ToNumberParser import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType, DatetimeType, Decimal, DecimalType, StringType} import org.apache.spark.unsafe.types.UTF8String @@ -50,7 +50,7 @@ abstract class ToNumberBase(left: Expression, right: Expression, errorOnFail: Bo } override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override def checkInputDataTypes(): TypeCheckResult = { val inputTypeCheck = super.checkInputDataTypes() @@ -284,7 +284,8 @@ case class ToCharacter(left: Expression, right: Expression) } override def dataType: DataType = SQLConf.get.defaultStringType - override def inputTypes: Seq[AbstractDataType] = Seq(DecimalType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(DecimalType, StringTypeWithCaseAccentSensitivity) override def checkInputDataTypes(): TypeCheckResult = { val inputTypeCheck = super.checkInputDataTypes() if (inputTypeCheck.isSuccess) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala index 970397c76a1cd..fdc3c27890469 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala @@ -35,7 +35,8 @@ import org.apache.spark.sql.catalyst.trees.BinaryLike import org.apache.spark.sql.catalyst.trees.TreePattern.{LIKE_FAMLIY, REGEXP_EXTRACT_FAMILY, REGEXP_REPLACE, TreePattern} import org.apache.spark.sql.catalyst.util.{CollationSupport, GenericArrayData, StringUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.internal.types.{StringTypeAnyCollation, StringTypeBinaryLcase} +import org.apache.spark.sql.internal.types.{ + StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -46,7 +47,7 @@ abstract class StringRegexExpression extends BinaryExpression def matches(regex: Pattern, str: String): Boolean override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity) final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId final lazy val collationRegexFlags: Int = CollationSupport.collationAwareRegexFlags(collationId) @@ -278,7 +279,7 @@ case class ILike( this(left, right, '\\') override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity) override protected def withNewChildrenInternal( newLeft: Expression, newRight: Expression): Expression = { @@ -567,7 +568,7 @@ case class StringSplit(str: Expression, regex: Expression, limit: Expression) override def dataType: DataType = ArrayType(str.dataType, containsNull = false) override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation, IntegerType) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity, IntegerType) override def first: Expression = str override def second: Expression = regex override def third: Expression = limit @@ -711,7 +712,8 @@ case class RegExpReplace(subject: Expression, regexp: Expression, rep: Expressio override def dataType: DataType = subject.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation, StringTypeBinaryLcase, IntegerType) + Seq(StringTypeBinaryLcase, + StringTypeWithCaseAccentSensitivity, StringTypeBinaryLcase, IntegerType) final lazy val collationId: Int = subject.dataType.asInstanceOf[StringType].collationId override def prettyName: String = "regexp_replace" @@ -799,7 +801,7 @@ abstract class RegExpExtractBase final override val nodePatterns: Seq[TreePattern] = Seq(REGEXP_EXTRACT_FAMILY) override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation, IntegerType) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity, IntegerType) override def first: Expression = subject override def second: Expression = regexp override def third: Expression = idx @@ -1052,7 +1054,7 @@ case class RegExpCount(left: Expression, right: Expression) override def children: Seq[Expression] = Seq(left, right) override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity) override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): RegExpCount = @@ -1092,7 +1094,7 @@ case class RegExpSubStr(left: Expression, right: Expression) override def children: Seq[Expression] = Seq(left, right) override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeBinaryLcase, StringTypeAnyCollation) + Seq(StringTypeBinaryLcase, StringTypeWithCaseAccentSensitivity) override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): RegExpSubStr = diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index 786c3968be0fe..c91c57ee1eb3e 100755 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -38,7 +38,8 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern, UPPER_OR_LO import org.apache.spark.sql.catalyst.util.{ArrayData, CharsetProvider, CollationFactory, CollationSupport, GenericArrayData, TypeUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeAnyCollation, StringTypeNonCSAICollation} +import org.apache.spark.sql.internal.types.{AbstractArrayType, + StringTypeNonCSAICollation, StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.UTF8StringBuilder import org.apache.spark.unsafe.array.ByteArrayMethods @@ -81,8 +82,10 @@ case class ConcatWs(children: Seq[Expression]) /** The 1st child (separator) is str, and rest are either str or array of str. */ override def inputTypes: Seq[AbstractDataType] = { val arrayOrStr = - TypeCollection(AbstractArrayType(StringTypeAnyCollation), StringTypeAnyCollation) - StringTypeAnyCollation +: Seq.fill(children.size - 1)(arrayOrStr) + TypeCollection(AbstractArrayType(StringTypeWithCaseAccentSensitivity), + StringTypeWithCaseAccentSensitivity + ) + StringTypeWithCaseAccentSensitivity +: Seq.fill(children.size - 1)(arrayOrStr) } override def dataType: DataType = children.head.dataType @@ -433,7 +436,7 @@ trait String2StringExpression extends ImplicitCastInputTypes { def convert(v: UTF8String): UTF8String override def dataType: DataType = child.dataType - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) protected override def nullSafeEval(input: Any): Any = convert(input.asInstanceOf[UTF8String]) @@ -515,7 +518,7 @@ abstract class StringPredicate extends BinaryExpression def compare(l: UTF8String, r: UTF8String): Boolean override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) protected override def nullSafeEval(input1: Any, input2: Any): Any = compare(input1.asInstanceOf[UTF8String], input2.asInstanceOf[UTF8String]) @@ -732,7 +735,7 @@ case class IsValidUTF8(input: Expression) extends RuntimeReplaceable with Implic override lazy val replacement: Expression = Invoke(input, "isValid", BooleanType) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nodeName: String = "is_valid_utf8" @@ -779,7 +782,7 @@ case class MakeValidUTF8(input: Expression) extends RuntimeReplaceable with Impl override lazy val replacement: Expression = Invoke(input, "makeValid", input.dataType) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nodeName: String = "make_valid_utf8" @@ -824,7 +827,7 @@ case class ValidateUTF8(input: Expression) extends RuntimeReplaceable with Impli Seq(input), inputTypes) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nodeName: String = "validate_utf8" @@ -873,7 +876,7 @@ case class TryValidateUTF8(input: Expression) extends RuntimeReplaceable with Im Seq(input), inputTypes) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nodeName: String = "try_validate_utf8" @@ -1008,8 +1011,8 @@ case class Overlay(input: Expression, replace: Expression, pos: Expression, len: override def dataType: DataType = input.dataType override def inputTypes: Seq[AbstractDataType] = Seq( - TypeCollection(StringTypeAnyCollation, BinaryType), - TypeCollection(StringTypeAnyCollation, BinaryType), IntegerType, IntegerType) + TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType), + TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType), IntegerType, IntegerType) override def checkInputDataTypes(): TypeCheckResult = { val inputTypeCheck = super.checkInputDataTypes() @@ -1213,7 +1216,7 @@ case class FindInSet(left: Expression, right: Expression) extends BinaryExpressi final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override protected def nullSafeEval(word: Any, set: Any): Any = { CollationSupport.FindInSet. @@ -1241,7 +1244,8 @@ trait String2TrimExpression extends Expression with ImplicitCastInputTypes { override def children: Seq[Expression] = srcStr +: trimStr.toSeq override def dataType: DataType = srcStr.dataType - override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq.fill(children.size)(StringTypeWithCaseAccentSensitivity) final lazy val collationId: Int = srcStr.dataType.asInstanceOf[StringType].collationId @@ -1846,7 +1850,7 @@ case class StringLPad(str: Expression, len: Expression, pad: Expression) override def dataType: DataType = str.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, IntegerType, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, IntegerType, StringTypeWithCaseAccentSensitivity) override def nullSafeEval(string: Any, len: Any, pad: Any): Any = { string.asInstanceOf[UTF8String].lpad(len.asInstanceOf[Int], pad.asInstanceOf[UTF8String]) @@ -1926,7 +1930,7 @@ case class StringRPad(str: Expression, len: Expression, pad: Expression = Litera override def dataType: DataType = str.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, IntegerType, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, IntegerType, StringTypeWithCaseAccentSensitivity) override def nullSafeEval(string: Any, len: Any, pad: Any): Any = { string.asInstanceOf[UTF8String].rpad(len.asInstanceOf[Int], pad.asInstanceOf[UTF8String]) @@ -1971,7 +1975,7 @@ case class FormatString(children: Expression*) extends Expression with ImplicitC override def dataType: DataType = children(0).dataType override def inputTypes: Seq[AbstractDataType] = - StringTypeAnyCollation :: List.fill(children.size - 1)(AnyDataType) + StringTypeWithCaseAccentSensitivity :: List.fill(children.size - 1)(AnyDataType) override def checkInputDataTypes(): TypeCheckResult = { if (children.isEmpty) { @@ -2082,7 +2086,7 @@ case class InitCap(child: Expression) // Flag to indicate whether to use ICU instead of JVM case mappings for UTF8_BINARY collation. private final lazy val useICU = SQLConf.get.getConf(SQLConf.ICU_CASE_MAPPINGS_ENABLED) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def dataType: DataType = child.dataType override def nullSafeEval(string: Any): Any = { @@ -2114,7 +2118,8 @@ case class StringRepeat(str: Expression, times: Expression) override def left: Expression = str override def right: Expression = times override def dataType: DataType = str.dataType - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation, IntegerType) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCaseAccentSensitivity, IntegerType) override def nullSafeEval(string: Any, n: Any): Any = { string.asInstanceOf[UTF8String].repeat(n.asInstanceOf[Integer]) @@ -2207,7 +2212,7 @@ case class Substring(str: Expression, pos: Expression, len: Expression) override def dataType: DataType = str.dataType override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, BinaryType), IntegerType, IntegerType) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType), IntegerType, IntegerType) override def first: Expression = str override def second: Expression = pos @@ -2265,7 +2270,8 @@ case class Right(str: Expression, len: Expression) extends RuntimeReplaceable ) ) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation, IntegerType) + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCaseAccentSensitivity, IntegerType) override def left: Expression = str override def right: Expression = len override protected def withNewChildrenInternal( @@ -2296,7 +2302,7 @@ case class Left(str: Expression, len: Expression) extends RuntimeReplaceable override lazy val replacement: Expression = Substring(str, Literal(1), len) override def inputTypes: Seq[AbstractDataType] = { - Seq(TypeCollection(StringTypeAnyCollation, BinaryType), IntegerType) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType), IntegerType) } override def left: Expression = str @@ -2332,7 +2338,7 @@ case class Length(child: Expression) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def dataType: DataType = IntegerType override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, BinaryType)) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType)) protected override def nullSafeEval(value: Any): Any = child.dataType match { case _: StringType => value.asInstanceOf[UTF8String].numChars @@ -2367,7 +2373,7 @@ case class BitLength(child: Expression) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def dataType: DataType = IntegerType override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, BinaryType)) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType)) protected override def nullSafeEval(value: Any): Any = child.dataType match { case _: StringType => value.asInstanceOf[UTF8String].numBytes * 8 @@ -2406,7 +2412,7 @@ case class OctetLength(child: Expression) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def dataType: DataType = IntegerType override def inputTypes: Seq[AbstractDataType] = - Seq(TypeCollection(StringTypeAnyCollation, BinaryType)) + Seq(TypeCollection(StringTypeWithCaseAccentSensitivity, BinaryType)) protected override def nullSafeEval(value: Any): Any = child.dataType match { case _: StringType => value.asInstanceOf[UTF8String].numBytes @@ -2466,8 +2472,9 @@ case class Levenshtein( } override def inputTypes: Seq[AbstractDataType] = threshold match { - case Some(_) => Seq(StringTypeAnyCollation, StringTypeAnyCollation, IntegerType) - case _ => Seq(StringTypeAnyCollation, StringTypeAnyCollation) + case Some(_) => + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity, IntegerType) + case _ => Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) } override def children: Seq[Expression] = threshold match { @@ -2592,7 +2599,7 @@ case class SoundEx(child: Expression) override def dataType: DataType = SQLConf.get.defaultStringType - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def nullSafeEval(input: Any): Any = input.asInstanceOf[UTF8String].soundex() @@ -2622,7 +2629,7 @@ case class Ascii(child: Expression) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def dataType: DataType = IntegerType - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) protected override def nullSafeEval(string: Any): Any = { // only pick the first character to reduce the `toString` cost @@ -2767,7 +2774,7 @@ case class UnBase64(child: Expression, failOnError: Boolean = false) extends UnaryExpression with ImplicitCastInputTypes with NullIntolerant { override def dataType: DataType = BinaryType - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) def this(expr: Expression) = this(expr, false) @@ -2946,7 +2953,8 @@ case class StringDecode( this(bin, charset, SQLConf.get.legacyJavaCharsets, SQLConf.get.legacyCodingErrorAction) override val dataType: DataType = SQLConf.get.defaultStringType - override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(BinaryType, StringTypeWithCaseAccentSensitivity) override def prettyName: String = "decode" override def toString: String = s"$prettyName($bin, $charset)" @@ -2955,7 +2963,7 @@ case class StringDecode( SQLConf.get.defaultStringType, "decode", Seq(bin, charset, Literal(legacyCharsets), Literal(legacyErrorAction)), - Seq(BinaryType, StringTypeAnyCollation, BooleanType, BooleanType)) + Seq(BinaryType, StringTypeWithCaseAccentSensitivity, BooleanType, BooleanType)) override def children: Seq[Expression] = Seq(bin, charset) override protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = @@ -3012,15 +3020,20 @@ case class Encode( override def dataType: DataType = BinaryType override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override lazy val replacement: Expression = StaticInvoke( classOf[Encode], BinaryType, "encode", Seq( - str, charset, Literal(legacyCharsets, BooleanType), Literal(legacyErrorAction, BooleanType)), - Seq(StringTypeAnyCollation, StringTypeAnyCollation, BooleanType, BooleanType)) + str, charset, Literal(legacyCharsets, BooleanType), Literal(legacyErrorAction, BooleanType) + ), + Seq( + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, + BooleanType, + BooleanType)) override def toString: String = s"$prettyName($str, $charset)" @@ -3104,7 +3117,8 @@ case class ToBinary( override def children: Seq[Expression] = expr +: format.toSeq - override def inputTypes: Seq[AbstractDataType] = children.map(_ => StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + children.map(_ => StringTypeWithCaseAccentSensitivity) override def checkInputDataTypes(): TypeCheckResult = { def isValidFormat: Boolean = { @@ -3120,7 +3134,8 @@ case class ToBinary( errorSubClass = "INVALID_ARG_VALUE", messageParameters = Map( "inputName" -> "fmt", - "requireType" -> s"case-insensitive ${toSQLType(StringTypeAnyCollation)}", + "requireType" -> + s"case-insensitive ${toSQLType(StringTypeWithCaseAccentSensitivity)}", "validValues" -> "'hex', 'utf-8', 'utf8', or 'base64'", "inputValue" -> toSQLValue(fmt, f.dataType) ) @@ -3131,7 +3146,7 @@ case class ToBinary( errorSubClass = "NON_FOLDABLE_INPUT", messageParameters = Map( "inputName" -> toSQLId("fmt"), - "inputType" -> toSQLType(StringTypeAnyCollation), + "inputType" -> toSQLType(StringTypeWithCaseAccentSensitivity), "inputExpr" -> toSQLExpr(f) ) ) @@ -3140,7 +3155,8 @@ case class ToBinary( errorSubClass = "INVALID_ARG_VALUE", messageParameters = Map( "inputName" -> "fmt", - "requireType" -> s"case-insensitive ${toSQLType(StringTypeAnyCollation)}", + "requireType" -> + s"case-insensitive ${toSQLType(StringTypeWithCaseAccentSensitivity)}", "validValues" -> "'hex', 'utf-8', 'utf8', or 'base64'", "inputValue" -> toSQLValue(f.eval(), f.dataType) ) @@ -3189,7 +3205,7 @@ case class FormatNumber(x: Expression, d: Expression) override def dataType: DataType = SQLConf.get.defaultStringType override def nullable: Boolean = true override def inputTypes: Seq[AbstractDataType] = - Seq(NumericType, TypeCollection(IntegerType, StringTypeAnyCollation)) + Seq(NumericType, TypeCollection(IntegerType, StringTypeWithCaseAccentSensitivity)) private val defaultFormat = "#,###,###,###,###,###,##0" @@ -3394,7 +3410,9 @@ case class Sentences( override def dataType: DataType = ArrayType(ArrayType(str.dataType, containsNull = false), containsNull = false) override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation, StringTypeAnyCollation) + Seq( + StringTypeWithCaseAccentSensitivity, + StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override def first: Expression = str override def second: Expression = language override def third: Expression = country @@ -3540,10 +3558,9 @@ case class Luhncheck(input: Expression) extends RuntimeReplaceable with Implicit classOf[ExpressionImplUtils], BooleanType, "isLuhnNumber", - Seq(input), - inputTypes) + Seq(input), inputTypes) - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def prettyName: String = "luhn_check" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/urlExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/urlExpressions.scala index 3e4e4f992002a..09e91da65484f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/urlExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/urlExpressions.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke import org.apache.spark.sql.catalyst.trees.UnaryLike import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types.{AbstractDataType, BooleanType, DataType} import org.apache.spark.unsafe.types.UTF8String @@ -59,13 +59,13 @@ case class UrlEncode(child: Expression) SQLConf.get.defaultStringType, "encode", Seq(child), - Seq(StringTypeAnyCollation)) + Seq(StringTypeWithCaseAccentSensitivity)) override protected def withNewChildInternal(newChild: Expression): Expression = { copy(child = newChild) } - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def prettyName: String = "url_encode" } @@ -98,13 +98,13 @@ case class UrlDecode(child: Expression, failOnError: Boolean = true) SQLConf.get.defaultStringType, "decode", Seq(child, Literal(failOnError)), - Seq(StringTypeAnyCollation, BooleanType)) + Seq(StringTypeWithCaseAccentSensitivity, BooleanType)) override protected def withNewChildInternal(newChild: Expression): Expression = { copy(child = newChild) } - override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCaseAccentSensitivity) override def prettyName: String = "url_decode" } @@ -190,7 +190,8 @@ case class ParseUrl(children: Seq[Expression], failOnError: Boolean = SQLConf.ge def this(children: Seq[Expression]) = this(children, SQLConf.get.ansiEnabled) override def nullable: Boolean = true - override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq.fill(children.size)(StringTypeWithCaseAccentSensitivity) override def dataType: DataType = SQLConf.get.defaultStringType override def prettyName: String = "parse_url" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala index 2c8ca1e8bb2bb..323f6e42f3e50 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData, import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.types.variant._ import org.apache.spark.types.variant.VariantUtil.{IntervalFields, Type} @@ -66,7 +66,7 @@ case class ParseJson(child: Expression, failOnError: Boolean = true) inputTypes :+ BooleanType :+ BooleanType, returnNullable = !failOnError) - override def inputTypes: Seq[AbstractDataType] = StringTypeAnyCollation :: Nil + override def inputTypes: Seq[AbstractDataType] = StringTypeWithCaseAccentSensitivity :: Nil override def dataType: DataType = VariantType @@ -271,7 +271,8 @@ case class VariantGet( final override def nodePatternsInternal(): Seq[TreePattern] = Seq(VARIANT_GET) - override def inputTypes: Seq[AbstractDataType] = Seq(VariantType, StringTypeAnyCollation) + override def inputTypes: Seq[AbstractDataType] = + Seq(VariantType, StringTypeWithCaseAccentSensitivity) override def prettyName: String = if (failOnError) "variant_get" else "try_variant_get" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala index 31e65cf0abc95..6c38bd88144b1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xml/xpath.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.Cast._ import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -42,7 +42,7 @@ abstract class XPathExtract override def nullable: Boolean = true override def inputTypes: Seq[AbstractDataType] = - Seq(StringTypeAnyCollation, StringTypeAnyCollation) + Seq(StringTypeWithCaseAccentSensitivity, StringTypeWithCaseAccentSensitivity) override def checkInputDataTypes(): TypeCheckResult = { if (!path.foldable) { @@ -50,7 +50,7 @@ abstract class XPathExtract errorSubClass = "NON_FOLDABLE_INPUT", messageParameters = Map( "inputName" -> toSQLId("path"), - "inputType" -> toSQLType(StringTypeAnyCollation), + "inputType" -> toSQLType(StringTypeWithCaseAccentSensitivity), "inputExpr" -> toSQLExpr(path) ) ) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala index 48a87db291a8d..6f1430b04ed67 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.util.TypeUtils._ import org.apache.spark.sql.catalyst.xml.{StaxXmlGenerator, StaxXmlParser, ValidatorUtil, XmlInferSchema, XmlOptions} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -124,7 +124,7 @@ case class XmlToStructs( defineCodeGen(ctx, ev, input => s"(InternalRow) $expr.nullSafeEval($input)") } - override def inputTypes: Seq[AbstractDataType] = StringTypeAnyCollation :: Nil + override def inputTypes: Seq[AbstractDataType] = StringTypeWithCaseAccentSensitivity :: Nil override def prettyName: String = "from_xml" diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala index de600d881b624..342dcbd8e6b6d 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeAnyCollation} +import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeWithCaseAccentSensitivity} import org.apache.spark.sql.types._ class AnsiTypeCoercionSuite extends TypeCoercionSuiteBase { @@ -1057,11 +1057,11 @@ class AnsiTypeCoercionSuite extends TypeCoercionSuiteBase { ArrayType(IntegerType)) shouldCast( ArrayType(StringType), - AbstractArrayType(StringTypeAnyCollation), + AbstractArrayType(StringTypeWithCaseAccentSensitivity), ArrayType(StringType)) shouldCast( ArrayType(IntegerType), - AbstractArrayType(StringTypeAnyCollation), + AbstractArrayType(StringTypeWithCaseAccentSensitivity), ArrayType(StringType)) shouldCast( ArrayType(StringType), @@ -1075,11 +1075,11 @@ class AnsiTypeCoercionSuite extends TypeCoercionSuiteBase { ArrayType(ArrayType(IntegerType))) shouldCast( ArrayType(ArrayType(StringType)), - AbstractArrayType(AbstractArrayType(StringTypeAnyCollation)), + AbstractArrayType(AbstractArrayType(StringTypeWithCaseAccentSensitivity)), ArrayType(ArrayType(StringType))) shouldCast( ArrayType(ArrayType(IntegerType)), - AbstractArrayType(AbstractArrayType(StringTypeAnyCollation)), + AbstractArrayType(AbstractArrayType(StringTypeWithCaseAccentSensitivity)), ArrayType(ArrayType(StringType))) shouldCast( ArrayType(ArrayType(StringType)), @@ -1088,14 +1088,16 @@ class AnsiTypeCoercionSuite extends TypeCoercionSuiteBase { // Invalid casts involving casting arrays into non-complex types. shouldNotCast(ArrayType(IntegerType), IntegerType) - shouldNotCast(ArrayType(StringType), StringTypeAnyCollation) + shouldNotCast(ArrayType(StringType), StringTypeWithCaseAccentSensitivity) shouldNotCast(ArrayType(StringType), IntegerType) - shouldNotCast(ArrayType(IntegerType), StringTypeAnyCollation) + shouldNotCast(ArrayType(IntegerType), StringTypeWithCaseAccentSensitivity) // Invalid casts involving casting arrays of arrays into arrays of non-complex types. shouldNotCast(ArrayType(ArrayType(IntegerType)), AbstractArrayType(IntegerType)) - shouldNotCast(ArrayType(ArrayType(StringType)), AbstractArrayType(StringTypeAnyCollation)) + shouldNotCast(ArrayType(ArrayType(StringType)), + AbstractArrayType(StringTypeWithCaseAccentSensitivity)) shouldNotCast(ArrayType(ArrayType(StringType)), AbstractArrayType(IntegerType)) - shouldNotCast(ArrayType(ArrayType(IntegerType)), AbstractArrayType(StringTypeAnyCollation)) + shouldNotCast(ArrayType(ArrayType(IntegerType)), + AbstractArrayType(StringTypeWithCaseAccentSensitivity)) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala index 9b454ba764f92..1aae2f10b7326 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjectio import org.apache.spark.sql.catalyst.util.CharsetProvider import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLId import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.StringTypeAnyCollation +import org.apache.spark.sql.internal.types.StringTypeWithCaseAccentSensitivity import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -1466,7 +1466,7 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { errorSubClass = "NON_FOLDABLE_INPUT", messageParameters = Map( "inputName" -> toSQLId("fmt"), - "inputType" -> toSQLType(StringTypeAnyCollation), + "inputType" -> toSQLType(StringTypeWithCaseAccentSensitivity), "inputExpr" -> toSQLExpr(wrongFmt) ) ) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala index 1d23774a51692..879c0c480943d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationExpressionWalkerSuite.scala @@ -66,10 +66,10 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi collationType: CollationType): Any = inputEntry match { case e: Class[_] if e.isAssignableFrom(classOf[Expression]) => - generateLiterals(StringTypeAnyCollation, collationType) + generateLiterals(StringTypeWithCaseAccentSensitivity, collationType) case se: Class[_] if se.isAssignableFrom(classOf[Seq[Expression]]) => - CreateArray(Seq(generateLiterals(StringTypeAnyCollation, collationType), - generateLiterals(StringTypeAnyCollation, collationType))) + CreateArray(Seq(generateLiterals(StringTypeWithCaseAccentSensitivity, collationType), + generateLiterals(StringTypeWithCaseAccentSensitivity, collationType))) case oe: Class[_] if oe.isAssignableFrom(classOf[Option[Any]]) => None case b: Class[_] if b.isAssignableFrom(classOf[Boolean]) => false case dt: Class[_] if dt.isAssignableFrom(classOf[DataType]) => StringType @@ -142,12 +142,12 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi lit => Literal.create(Seq(lit.asInstanceOf[Literal].value), ArrayType(lit.dataType)) ).head case ArrayType => - generateLiterals(StringTypeAnyCollation, collationType).map( + generateLiterals(StringTypeWithCaseAccentSensitivity, collationType).map( lit => Literal.create(Seq(lit.asInstanceOf[Literal].value), ArrayType(lit.dataType)) ).head case MapType => - val key = generateLiterals(StringTypeAnyCollation, collationType) - val value = generateLiterals(StringTypeAnyCollation, collationType) + val key = generateLiterals(StringTypeWithCaseAccentSensitivity, collationType) + val value = generateLiterals(StringTypeWithCaseAccentSensitivity, collationType) CreateMap(Seq(key, value)) case MapType(keyType, valueType, _) => val key = generateLiterals(keyType, collationType) @@ -159,8 +159,9 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi CreateMap(Seq(key, value)) case StructType => CreateNamedStruct( - Seq(Literal("start"), generateLiterals(StringTypeAnyCollation, collationType), - Literal("end"), generateLiterals(StringTypeAnyCollation, collationType))) + Seq(Literal("start"), + generateLiterals(StringTypeWithCaseAccentSensitivity, collationType), + Literal("end"), generateLiterals(StringTypeWithCaseAccentSensitivity, collationType))) } /** @@ -209,10 +210,10 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case ArrayType(elementType, _) => "array(" + generateInputAsString(elementType, collationType) + ")" case ArrayType => - "array(" + generateInputAsString(StringTypeAnyCollation, collationType) + ")" + "array(" + generateInputAsString(StringTypeWithCaseAccentSensitivity, collationType) + ")" case MapType => - "map(" + generateInputAsString(StringTypeAnyCollation, collationType) + ", " + - generateInputAsString(StringTypeAnyCollation, collationType) + ")" + "map(" + generateInputAsString(StringTypeWithCaseAccentSensitivity, collationType) + ", " + + generateInputAsString(StringTypeWithCaseAccentSensitivity, collationType) + ")" case MapType(keyType, valueType, _) => "map(" + generateInputAsString(keyType, collationType) + ", " + generateInputAsString(valueType, collationType) + ")" @@ -220,8 +221,9 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi "map(" + generateInputAsString(keyType, collationType) + ", " + generateInputAsString(valueType, collationType) + ")" case StructType => - "named_struct( 'start', " + generateInputAsString(StringTypeAnyCollation, collationType) + - ", 'end', " + generateInputAsString(StringTypeAnyCollation, collationType) + ")" + "named_struct( 'start', " + + generateInputAsString(StringTypeWithCaseAccentSensitivity, collationType) + ", 'end', " + + generateInputAsString(StringTypeWithCaseAccentSensitivity, collationType) + ")" case StructType(fields) => "named_struct(" + fields.map(f => "'" + f.name + "', " + generateInputAsString(f.dataType, collationType)).mkString(", ") + ")" @@ -267,10 +269,12 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi case ArrayType(elementType, _) => "array<" + generateInputTypeAsStrings(elementType, collationType) + ">" case ArrayType => - "array<" + generateInputTypeAsStrings(StringTypeAnyCollation, collationType) + ">" + "array<" + generateInputTypeAsStrings(StringTypeWithCaseAccentSensitivity, collationType) + + ">" case MapType => - "map<" + generateInputTypeAsStrings(StringTypeAnyCollation, collationType) + ", " + - generateInputTypeAsStrings(StringTypeAnyCollation, collationType) + ">" + "map<" + generateInputTypeAsStrings(StringTypeWithCaseAccentSensitivity, collationType) + + ", " + + generateInputTypeAsStrings(StringTypeWithCaseAccentSensitivity, collationType) + ">" case MapType(keyType, valueType, _) => "map<" + generateInputTypeAsStrings(keyType, collationType) + ", " + generateInputTypeAsStrings(valueType, collationType) + ">" @@ -278,9 +282,10 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi "map<" + generateInputTypeAsStrings(keyType, collationType) + ", " + generateInputTypeAsStrings(valueType, collationType) + ">" case StructType => - "struct" + generateInputTypeAsStrings(StringTypeWithCaseAccentSensitivity, collationType) + ">" case StructType(fields) => "named_struct<" + fields.map(f => "'" + f.name + "', " + generateInputTypeAsStrings(f.dataType, collationType)).mkString(", ") + ">" @@ -293,8 +298,8 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi */ def hasStringType(inputType: AbstractDataType): Boolean = { inputType match { - case _: StringType | StringTypeAnyCollation | StringTypeBinaryLcase | AnyDataType => - true + case _: StringType | StringTypeWithCaseAccentSensitivity | StringTypeBinaryLcase | AnyDataType + => true case ArrayType => true case MapType => true case MapType(keyType, valueType, _) => hasStringType(keyType) || hasStringType(valueType) @@ -408,7 +413,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi var input: Seq[Expression] = Seq.empty var i = 0 for (_ <- 1 to 10) { - input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + input = input :+ generateLiterals(StringTypeWithCaseAccentSensitivity, Utf8Binary) try { method.invoke(null, funInfo.getClassName, input).asInstanceOf[ExpectsInputTypes] } @@ -498,7 +503,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi var input: Seq[Expression] = Seq.empty var result: Expression = null for (_ <- 1 to 10) { - input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + input = input :+ generateLiterals(StringTypeWithCaseAccentSensitivity, Utf8Binary) try { val tempResult = method.invoke(null, f.getClassName, input) if (result == null) result = tempResult.asInstanceOf[Expression] @@ -609,7 +614,7 @@ class CollationExpressionWalkerSuite extends SparkFunSuite with SharedSparkSessi var input: Seq[Expression] = Seq.empty var result: Expression = null for (_ <- 1 to 10) { - input = input :+ generateLiterals(StringTypeAnyCollation, Utf8Binary) + input = input :+ generateLiterals(StringTypeWithCaseAccentSensitivity, Utf8Binary) try { val tempResult = method.invoke(null, f.getClassName, input) if (result == null) result = tempResult.asInstanceOf[Expression] From c54c017e93090a5fb2edf1b5ef029561b6387a3f Mon Sep 17 00:00:00 2001 From: Jovan Pavlovic Date: Mon, 30 Sep 2024 17:44:13 +0800 Subject: [PATCH 174/189] [SPARK-49666][SQL] Add feature flag for trim collation feature ### What changes were proposed in this pull request? Introducing new specifier for trim collations (both leading and trailing trimming). These are initial changes so that trim specifier is recognized and put under feature flag (all code paths blocked). ### Why are the changes needed? Support for trailing space trimming is one of the requested feature by users. ### Does this PR introduce _any_ user-facing change? This is guarded by feature flag. ### How was this patch tested? Added tests to CollationSuite, SqlConfSuite and QueryCompilationErrorSuite. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48222 from jovanpavl-db/trim-collation-feature-initial-support. Authored-by: Jovan Pavlovic Signed-off-by: Wenchen Fan --- .../sql/catalyst/util/CollationFactory.java | 341 +++++++++++++----- .../unsafe/types/CollationFactorySuite.scala | 5 +- .../resources/error/error-conditions.json | 10 +- .../expressions/collationExpressions.scala | 4 + .../sql/catalyst/parser/AstBuilder.scala | 4 + .../sql/errors/QueryCompilationErrors.scala | 7 + .../apache/spark/sql/internal/SQLConf.scala | 14 + .../spark/sql/execution/SparkSqlParser.scala | 4 + .../org/apache/spark/sql/CollationSuite.scala | 56 ++- .../errors/QueryCompilationErrorsSuite.scala | 33 ++ .../spark/sql/internal/SQLConfSuite.scala | 7 + 11 files changed, 381 insertions(+), 104 deletions(-) diff --git a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java index d5dbca7eb89bc..e368e2479a3a1 100644 --- a/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java +++ b/common/unsafe/src/main/java/org/apache/spark/sql/catalyst/util/CollationFactory.java @@ -99,7 +99,8 @@ public record CollationMeta( String icuVersion, String padAttribute, boolean accentSensitivity, - boolean caseSensitivity) { } + boolean caseSensitivity, + String spaceTrimming) { } /** * Entry encapsulating all information about a collation. @@ -200,6 +201,7 @@ public Collation( * bit 28-24: Reserved. * bit 23-22: Reserved for version. * bit 21-18: Reserved for space trimming. + * 0000 = none, 0001 = left trim, 0010 = right trim, 0011 = trim. * bit 17-0: Depend on collation family. * --- * INDETERMINATE collation ID binary layout: @@ -214,7 +216,8 @@ public Collation( * UTF8_BINARY collation ID binary layout: * bit 31-24: Zeroes. * bit 23-22: Zeroes, reserved for version. - * bit 21-18: Zeroes, reserved for space trimming. + * bit 21-18: Reserved for space trimming. + * 0000 = none, 0001 = left trim, 0010 = right trim, 0011 = trim. * bit 17-3: Zeroes. * bit 2: 0, reserved for accent sensitivity. * bit 1: 0, reserved for uppercase and case-insensitive. @@ -225,7 +228,8 @@ public Collation( * bit 29: 1 * bit 28-24: Zeroes. * bit 23-22: Zeroes, reserved for version. - * bit 21-18: Zeroes, reserved for space trimming. + * bit 21-18: Reserved for space trimming. + * 0000 = none, 0001 = left trim, 0010 = right trim, 0011 = trim. * bit 17: 0 = case-sensitive, 1 = case-insensitive. * bit 16: 0 = accent-sensitive, 1 = accent-insensitive. * bit 15-14: Zeroes, reserved for punctuation sensitivity. @@ -238,7 +242,13 @@ public Collation( * - UNICODE -> 0x20000000 * - UNICODE_AI -> 0x20010000 * - UNICODE_CI -> 0x20020000 + * - UNICODE_LTRIM -> 0x20040000 + * - UNICODE_RTRIM -> 0x20080000 + * - UNICODE_TRIM -> 0x200C0000 * - UNICODE_CI_AI -> 0x20030000 + * - UNICODE_CI_TRIM -> 0x200E0000 + * - UNICODE_AI_TRIM -> 0x200D0000 + * - UNICODE_CI_AI_TRIM-> 0x200F0000 * - af -> 0x20000001 * - af_CI_AI -> 0x20030001 */ @@ -259,6 +269,15 @@ protected enum ImplementationProvider { UTF8_BINARY, ICU } + /** + * Bits 19-18 having value 00 for no space trimming, 01 for left space trimming + * 10 for right space trimming and 11 for both sides space trimming. Bits 21, 20 + * remained reserved (and fixed to 0) for future use. + */ + protected enum SpaceTrimming { + NONE, LTRIM, RTRIM, TRIM + } + /** * Offset in binary collation ID layout. */ @@ -279,6 +298,17 @@ protected enum ImplementationProvider { */ protected static final int IMPLEMENTATION_PROVIDER_MASK = 0b1; + + /** + * Offset in binary collation ID layout. + */ + protected static final int SPACE_TRIMMING_OFFSET = 18; + + /** + * Bitmask corresponding to width in bits in binary collation ID layout. + */ + protected static final int SPACE_TRIMMING_MASK = 0b11; + private static final int INDETERMINATE_COLLATION_ID = -1; /** @@ -303,6 +333,14 @@ private static DefinitionOrigin getDefinitionOrigin(int collationId) { DEFINITION_ORIGIN_OFFSET, DEFINITION_ORIGIN_MASK)]; } + /** + * Utility function to retrieve `SpaceTrimming` enum instance from collation ID. + */ + protected static SpaceTrimming getSpaceTrimming(int collationId) { + return SpaceTrimming.values()[SpecifierUtils.getSpecValue(collationId, + SPACE_TRIMMING_OFFSET, SPACE_TRIMMING_MASK)]; + } + /** * Main entry point for retrieving `Collation` instance from collation ID. */ @@ -358,6 +396,8 @@ private static int collationNameToId(String collationName) throws SparkException protected abstract CollationMeta buildCollationMeta(); + protected abstract String normalizedCollationName(); + static List listCollations() { return Stream.concat( CollationSpecUTF8.listCollations().stream(), @@ -398,48 +438,99 @@ private enum CaseSensitivity { private static final String UTF8_LCASE_COLLATION_NAME = "UTF8_LCASE"; private static final int UTF8_BINARY_COLLATION_ID = - new CollationSpecUTF8(CaseSensitivity.UNSPECIFIED).collationId; + new CollationSpecUTF8(CaseSensitivity.UNSPECIFIED, SpaceTrimming.NONE).collationId; private static final int UTF8_LCASE_COLLATION_ID = - new CollationSpecUTF8(CaseSensitivity.LCASE).collationId; + new CollationSpecUTF8(CaseSensitivity.LCASE, SpaceTrimming.NONE).collationId; protected static Collation UTF8_BINARY_COLLATION = - new CollationSpecUTF8(CaseSensitivity.UNSPECIFIED).buildCollation(); + new CollationSpecUTF8(CaseSensitivity.UNSPECIFIED, SpaceTrimming.NONE).buildCollation(); protected static Collation UTF8_LCASE_COLLATION = - new CollationSpecUTF8(CaseSensitivity.LCASE).buildCollation(); + new CollationSpecUTF8(CaseSensitivity.LCASE, SpaceTrimming.NONE).buildCollation(); + private final CaseSensitivity caseSensitivity; + private final SpaceTrimming spaceTrimming; private final int collationId; - private CollationSpecUTF8(CaseSensitivity caseSensitivity) { - this.collationId = + private CollationSpecUTF8( + CaseSensitivity caseSensitivity, + SpaceTrimming spaceTrimming) { + this.caseSensitivity = caseSensitivity; + this.spaceTrimming = spaceTrimming; + + int collationId = SpecifierUtils.setSpecValue(0, CASE_SENSITIVITY_OFFSET, caseSensitivity); + this.collationId = + SpecifierUtils.setSpecValue(collationId, SPACE_TRIMMING_OFFSET, spaceTrimming); } private static int collationNameToId(String originalName, String collationName) throws SparkException { - if (UTF8_BINARY_COLLATION.collationName.equals(collationName)) { - return UTF8_BINARY_COLLATION_ID; - } else if (UTF8_LCASE_COLLATION.collationName.equals(collationName)) { - return UTF8_LCASE_COLLATION_ID; + + int baseId; + String collationNamePrefix; + + if (collationName.startsWith(UTF8_BINARY_COLLATION.collationName)) { + baseId = UTF8_BINARY_COLLATION_ID; + collationNamePrefix = UTF8_BINARY_COLLATION.collationName; + } else if (collationName.startsWith(UTF8_LCASE_COLLATION.collationName)) { + baseId = UTF8_LCASE_COLLATION_ID; + collationNamePrefix = UTF8_LCASE_COLLATION.collationName; } else { // Throw exception with original (before case conversion) collation name. throw collationInvalidNameException(originalName); } + + String remainingSpecifiers = collationName.substring(collationNamePrefix.length()); + if(remainingSpecifiers.isEmpty()) { + return baseId; + } + if(!remainingSpecifiers.startsWith("_")){ + throw collationInvalidNameException(originalName); + } + + SpaceTrimming spaceTrimming = SpaceTrimming.NONE; + String remainingSpec = remainingSpecifiers.substring(1); + if (remainingSpec.equals("LTRIM")) { + spaceTrimming = SpaceTrimming.LTRIM; + } else if (remainingSpec.equals("RTRIM")) { + spaceTrimming = SpaceTrimming.RTRIM; + } else if(remainingSpec.equals("TRIM")) { + spaceTrimming = SpaceTrimming.TRIM; + } else { + throw collationInvalidNameException(originalName); + } + + return SpecifierUtils.setSpecValue(baseId, SPACE_TRIMMING_OFFSET, spaceTrimming); } private static CollationSpecUTF8 fromCollationId(int collationId) { // Extract case sensitivity from collation ID. int caseConversionOrdinal = SpecifierUtils.getSpecValue(collationId, CASE_SENSITIVITY_OFFSET, CASE_SENSITIVITY_MASK); - // Verify only case sensitivity bits were set settable in UTF8_BINARY family of collations. - assert (SpecifierUtils.removeSpec(collationId, - CASE_SENSITIVITY_OFFSET, CASE_SENSITIVITY_MASK) == 0); - return new CollationSpecUTF8(CaseSensitivity.values()[caseConversionOrdinal]); + // Extract space trimming from collation ID. + int spaceTrimmingOrdinal = getSpaceTrimming(collationId).ordinal(); + assert(isValidCollationId(collationId)); + return new CollationSpecUTF8( + CaseSensitivity.values()[caseConversionOrdinal], + SpaceTrimming.values()[spaceTrimmingOrdinal]); + } + + private static boolean isValidCollationId(int collationId) { + collationId = SpecifierUtils.removeSpec( + collationId, + SPACE_TRIMMING_OFFSET, + SPACE_TRIMMING_MASK); + collationId = SpecifierUtils.removeSpec( + collationId, + CASE_SENSITIVITY_OFFSET, + CASE_SENSITIVITY_MASK); + return collationId == 0; } @Override protected Collation buildCollation() { - if (collationId == UTF8_BINARY_COLLATION_ID) { + if (caseSensitivity == CaseSensitivity.UNSPECIFIED) { return new Collation( - UTF8_BINARY_COLLATION_NAME, + normalizedCollationName(), PROVIDER_SPARK, null, UTF8String::binaryCompare, @@ -450,7 +541,7 @@ protected Collation buildCollation() { /* supportsLowercaseEquality = */ false); } else { return new Collation( - UTF8_LCASE_COLLATION_NAME, + normalizedCollationName(), PROVIDER_SPARK, null, CollationAwareUTF8String::compareLowerCase, @@ -464,29 +555,52 @@ protected Collation buildCollation() { @Override protected CollationMeta buildCollationMeta() { - if (collationId == UTF8_BINARY_COLLATION_ID) { + if (caseSensitivity == CaseSensitivity.UNSPECIFIED) { return new CollationMeta( CATALOG, SCHEMA, - UTF8_BINARY_COLLATION_NAME, + normalizedCollationName(), /* language = */ null, /* country = */ null, /* icuVersion = */ null, COLLATION_PAD_ATTRIBUTE, /* accentSensitivity = */ true, - /* caseSensitivity = */ true); + /* caseSensitivity = */ true, + spaceTrimming.toString()); } else { return new CollationMeta( CATALOG, SCHEMA, - UTF8_LCASE_COLLATION_NAME, + normalizedCollationName(), /* language = */ null, /* country = */ null, /* icuVersion = */ null, COLLATION_PAD_ATTRIBUTE, /* accentSensitivity = */ true, - /* caseSensitivity = */ false); + /* caseSensitivity = */ false, + spaceTrimming.toString()); + } + } + + /** + * Compute normalized collation name. Components of collation name are given in order: + * - Base collation name (UTF8_BINARY or UTF8_LCASE) + * - Optional space trimming when non-default preceded by underscore + * Examples: UTF8_BINARY, UTF8_BINARY_LCASE_LTRIM, UTF8_BINARY_TRIM. + */ + @Override + protected String normalizedCollationName() { + StringBuilder builder = new StringBuilder(); + if(caseSensitivity == CaseSensitivity.UNSPECIFIED){ + builder.append(UTF8_BINARY_COLLATION_NAME); + } else{ + builder.append(UTF8_LCASE_COLLATION_NAME); } + if (spaceTrimming != SpaceTrimming.NONE) { + builder.append('_'); + builder.append(spaceTrimming.toString()); + } + return builder.toString(); } static List listCollations() { @@ -620,21 +734,33 @@ private enum AccentSensitivity { } } - private static final int UNICODE_COLLATION_ID = - new CollationSpecICU("UNICODE", CaseSensitivity.CS, AccentSensitivity.AS).collationId; - private static final int UNICODE_CI_COLLATION_ID = - new CollationSpecICU("UNICODE", CaseSensitivity.CI, AccentSensitivity.AS).collationId; + private static final int UNICODE_COLLATION_ID = new CollationSpecICU( + "UNICODE", + CaseSensitivity.CS, + AccentSensitivity.AS, + SpaceTrimming.NONE).collationId; + + private static final int UNICODE_CI_COLLATION_ID = new CollationSpecICU( + "UNICODE", + CaseSensitivity.CI, + AccentSensitivity.AS, + SpaceTrimming.NONE).collationId; private final CaseSensitivity caseSensitivity; private final AccentSensitivity accentSensitivity; + private final SpaceTrimming spaceTrimming; private final String locale; private final int collationId; - private CollationSpecICU(String locale, CaseSensitivity caseSensitivity, - AccentSensitivity accentSensitivity) { + private CollationSpecICU( + String locale, + CaseSensitivity caseSensitivity, + AccentSensitivity accentSensitivity, + SpaceTrimming spaceTrimming) { this.locale = locale; this.caseSensitivity = caseSensitivity; this.accentSensitivity = accentSensitivity; + this.spaceTrimming = spaceTrimming; // Construct collation ID from locale, case-sensitivity and accent-sensitivity specifiers. int collationId = ICULocaleToId.get(locale); // Mandatory ICU implementation provider. @@ -644,6 +770,8 @@ private CollationSpecICU(String locale, CaseSensitivity caseSensitivity, caseSensitivity); collationId = SpecifierUtils.setSpecValue(collationId, ACCENT_SENSITIVITY_OFFSET, accentSensitivity); + collationId = SpecifierUtils.setSpecValue(collationId, SPACE_TRIMMING_OFFSET, + spaceTrimming); this.collationId = collationId; } @@ -661,58 +789,88 @@ private static int collationNameToId( } if (lastPos == -1) { throw collationInvalidNameException(originalName); - } else { - String locale = collationName.substring(0, lastPos); - int collationId = ICULocaleToId.get(ICULocaleMapUppercase.get(locale)); - - // Try all combinations of AS/AI and CS/CI. - CaseSensitivity caseSensitivity; - AccentSensitivity accentSensitivity; - if (collationName.equals(locale) || - collationName.equals(locale + "_AS") || - collationName.equals(locale + "_CS") || - collationName.equals(locale + "_AS_CS") || - collationName.equals(locale + "_CS_AS") - ) { - caseSensitivity = CaseSensitivity.CS; - accentSensitivity = AccentSensitivity.AS; - } else if (collationName.equals(locale + "_CI") || - collationName.equals(locale + "_AS_CI") || - collationName.equals(locale + "_CI_AS")) { - caseSensitivity = CaseSensitivity.CI; - accentSensitivity = AccentSensitivity.AS; - } else if (collationName.equals(locale + "_AI") || - collationName.equals(locale + "_CS_AI") || - collationName.equals(locale + "_AI_CS")) { - caseSensitivity = CaseSensitivity.CS; - accentSensitivity = AccentSensitivity.AI; - } else if (collationName.equals(locale + "_AI_CI") || - collationName.equals(locale + "_CI_AI")) { - caseSensitivity = CaseSensitivity.CI; - accentSensitivity = AccentSensitivity.AI; - } else { - throw collationInvalidNameException(originalName); - } + } + String locale = collationName.substring(0, lastPos); + int collationId = ICULocaleToId.get(ICULocaleMapUppercase.get(locale)); + collationId = SpecifierUtils.setSpecValue(collationId, + IMPLEMENTATION_PROVIDER_OFFSET, ImplementationProvider.ICU); - // Build collation ID from computed specifiers. - collationId = SpecifierUtils.setSpecValue(collationId, - IMPLEMENTATION_PROVIDER_OFFSET, ImplementationProvider.ICU); - collationId = SpecifierUtils.setSpecValue(collationId, - CASE_SENSITIVITY_OFFSET, caseSensitivity); - collationId = SpecifierUtils.setSpecValue(collationId, - ACCENT_SENSITIVITY_OFFSET, accentSensitivity); + // No other specifiers present. + if(collationName.equals(locale)){ return collationId; } + if(collationName.charAt(locale.length()) != '_'){ + throw collationInvalidNameException(originalName); + } + // Extract remaining specifiers and trim "_" separator. + String remainingSpecifiers = collationName.substring(lastPos + 1); + + // Initialize default specifier flags. + // Case sensitive, accent sensitive, no space trimming. + boolean isCaseSpecifierSet = false; + boolean isAccentSpecifierSet = false; + boolean isSpaceTrimmingSpecifierSet = false; + CaseSensitivity caseSensitivity = CaseSensitivity.CS; + AccentSensitivity accentSensitivity = AccentSensitivity.AS; + SpaceTrimming spaceTrimming = SpaceTrimming.NONE; + + String[] specifiers = remainingSpecifiers.split("_"); + + // Iterate through specifiers and set corresponding flags + for (String specifier : specifiers) { + switch (specifier) { + case "CI": + case "CS": + if (isCaseSpecifierSet) { + throw collationInvalidNameException(originalName); + } + caseSensitivity = CaseSensitivity.valueOf(specifier); + isCaseSpecifierSet = true; + break; + case "AI": + case "AS": + if (isAccentSpecifierSet) { + throw collationInvalidNameException(originalName); + } + accentSensitivity = AccentSensitivity.valueOf(specifier); + isAccentSpecifierSet = true; + break; + case "LTRIM": + case "RTRIM": + case "TRIM": + if (isSpaceTrimmingSpecifierSet) { + throw collationInvalidNameException(originalName); + } + spaceTrimming = SpaceTrimming.valueOf(specifier); + isSpaceTrimmingSpecifierSet = true; + break; + default: + throw collationInvalidNameException(originalName); + } + } + + // Build collation ID from computed specifiers. + collationId = SpecifierUtils.setSpecValue(collationId, + CASE_SENSITIVITY_OFFSET, caseSensitivity); + collationId = SpecifierUtils.setSpecValue(collationId, + ACCENT_SENSITIVITY_OFFSET, accentSensitivity); + collationId = SpecifierUtils.setSpecValue(collationId, + SPACE_TRIMMING_OFFSET, spaceTrimming); + return collationId; } private static CollationSpecICU fromCollationId(int collationId) { // Parse specifiers from collation ID. + int spaceTrimmingOrdinal = SpecifierUtils.getSpecValue(collationId, + SPACE_TRIMMING_OFFSET, SPACE_TRIMMING_MASK); int caseSensitivityOrdinal = SpecifierUtils.getSpecValue(collationId, CASE_SENSITIVITY_OFFSET, CASE_SENSITIVITY_MASK); int accentSensitivityOrdinal = SpecifierUtils.getSpecValue(collationId, ACCENT_SENSITIVITY_OFFSET, ACCENT_SENSITIVITY_MASK); collationId = SpecifierUtils.removeSpec(collationId, IMPLEMENTATION_PROVIDER_OFFSET, IMPLEMENTATION_PROVIDER_MASK); + collationId = SpecifierUtils.removeSpec(collationId, + SPACE_TRIMMING_OFFSET, SPACE_TRIMMING_MASK); collationId = SpecifierUtils.removeSpec(collationId, CASE_SENSITIVITY_OFFSET, CASE_SENSITIVITY_MASK); collationId = SpecifierUtils.removeSpec(collationId, @@ -723,8 +881,9 @@ private static CollationSpecICU fromCollationId(int collationId) { assert(localeId >= 0 && localeId < ICULocaleNames.length); CaseSensitivity caseSensitivity = CaseSensitivity.values()[caseSensitivityOrdinal]; AccentSensitivity accentSensitivity = AccentSensitivity.values()[accentSensitivityOrdinal]; + SpaceTrimming spaceTrimming = SpaceTrimming.values()[spaceTrimmingOrdinal]; String locale = ICULocaleNames[localeId]; - return new CollationSpecICU(locale, caseSensitivity, accentSensitivity); + return new CollationSpecICU(locale, caseSensitivity, accentSensitivity, spaceTrimming); } @Override @@ -752,7 +911,7 @@ protected Collation buildCollation() { // Freeze ICU collator to ensure thread safety. collator.freeze(); return new Collation( - collationName(), + normalizedCollationName(), PROVIDER_ICU, collator, (s1, s2) -> collator.compare(s1.toValidString(), s2.toValidString()), @@ -768,13 +927,14 @@ protected CollationMeta buildCollationMeta() { return new CollationMeta( CATALOG, SCHEMA, - collationName(), + normalizedCollationName(), ICULocaleMap.get(locale).getDisplayLanguage(), ICULocaleMap.get(locale).getDisplayCountry(), VersionInfo.ICU_VERSION.toString(), COLLATION_PAD_ATTRIBUTE, accentSensitivity == AccentSensitivity.AS, - caseSensitivity == CaseSensitivity.CS); + caseSensitivity == CaseSensitivity.CS, + spaceTrimming.toString()); } /** @@ -782,9 +942,11 @@ protected CollationMeta buildCollationMeta() { * - Locale name * - Optional case sensitivity when non-default preceded by underscore * - Optional accent sensitivity when non-default preceded by underscore - * Examples: en, en_USA_CI_AI, sr_Cyrl_SRB_AI. + * - Optional space trimming when non-default preceded by underscore + * Examples: en, en_USA_CI_LTRIM, en_USA_CI_AI, en_USA_CI_AI_TRIM, sr_Cyrl_SRB_AI. */ - private String collationName() { + @Override + protected String normalizedCollationName() { StringBuilder builder = new StringBuilder(); builder.append(locale); if (caseSensitivity != CaseSensitivity.CS) { @@ -795,20 +957,21 @@ private String collationName() { builder.append('_'); builder.append(accentSensitivity.toString()); } + if(spaceTrimming != SpaceTrimming.NONE) { + builder.append('_'); + builder.append(spaceTrimming.toString()); + } return builder.toString(); } private static List allCollationNames() { List collationNames = new ArrayList<>(); - for (String locale: ICULocaleToId.keySet()) { - // CaseSensitivity.CS + AccentSensitivity.AS - collationNames.add(locale); - // CaseSensitivity.CS + AccentSensitivity.AI - collationNames.add(locale + "_AI"); - // CaseSensitivity.CI + AccentSensitivity.AS - collationNames.add(locale + "_CI"); - // CaseSensitivity.CI + AccentSensitivity.AI - collationNames.add(locale + "_CI_AI"); + List caseAccentSpecifiers = Arrays.asList("", "_AI", "_CI", "_CI_AI"); + for (String locale : ICULocaleToId.keySet()) { + for (String caseAccent : caseAccentSpecifiers) { + String collationName = locale + caseAccent; + collationNames.add(collationName); + } } return collationNames.stream().sorted().toList(); } @@ -933,6 +1096,14 @@ public static boolean isCaseSensitiveAndAccentInsensitive(int collationId) { Collation.CollationSpecICU.AccentSensitivity.AI; } + /** + * Returns whether the collation uses trim collation for the given collation id. + */ + public static boolean usesTrimCollation(int collationId) { + return Collation.CollationSpec.getSpaceTrimming(collationId) != + Collation.CollationSpec.SpaceTrimming.NONE; + } + public static void assertValidProvider(String provider) throws SparkException { if (!SUPPORTED_PROVIDERS.contains(provider.toLowerCase())) { Map params = Map.of( diff --git a/common/unsafe/src/test/scala/org/apache/spark/unsafe/types/CollationFactorySuite.scala b/common/unsafe/src/test/scala/org/apache/spark/unsafe/types/CollationFactorySuite.scala index 321d1ccd700f2..054c44f7286b7 100644 --- a/common/unsafe/src/test/scala/org/apache/spark/unsafe/types/CollationFactorySuite.scala +++ b/common/unsafe/src/test/scala/org/apache/spark/unsafe/types/CollationFactorySuite.scala @@ -369,9 +369,8 @@ class CollationFactorySuite extends AnyFunSuite with Matchers { // scalastyle:ig 1 << 15, // UTF8_BINARY mandatory zero bit 15 breach. 1 << 16, // UTF8_BINARY mandatory zero bit 16 breach. 1 << 17, // UTF8_BINARY mandatory zero bit 17 breach. - 1 << 18, // UTF8_BINARY mandatory zero bit 18 breach. - 1 << 19, // UTF8_BINARY mandatory zero bit 19 breach. 1 << 20, // UTF8_BINARY mandatory zero bit 20 breach. + 1 << 21, // UTF8_BINARY mandatory zero bit 21 breach. 1 << 23, // UTF8_BINARY mandatory zero bit 23 breach. 1 << 24, // UTF8_BINARY mandatory zero bit 24 breach. 1 << 25, // UTF8_BINARY mandatory zero bit 25 breach. @@ -382,8 +381,6 @@ class CollationFactorySuite extends AnyFunSuite with Matchers { // scalastyle:ig (1 << 29) | (1 << 13), // ICU mandatory zero bit 13 breach. (1 << 29) | (1 << 14), // ICU mandatory zero bit 14 breach. (1 << 29) | (1 << 15), // ICU mandatory zero bit 15 breach. - (1 << 29) | (1 << 18), // ICU mandatory zero bit 18 breach. - (1 << 29) | (1 << 19), // ICU mandatory zero bit 19 breach. (1 << 29) | (1 << 20), // ICU mandatory zero bit 20 breach. (1 << 29) | (1 << 21), // ICU mandatory zero bit 21 breach. (1 << 29) | (1 << 22), // ICU mandatory zero bit 22 breach. diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 3fcb53426eccf..fcaf2b1d9d301 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4886,11 +4886,6 @@ "Catalog does not support ." ] }, - "COLLATION" : { - "message" : [ - "Collation is not yet supported." - ] - }, "COMBINATION_QUERY_RESULT_CLAUSES" : { "message" : [ "Combination of ORDER BY/SORT BY/DISTRIBUTE BY/CLUSTER BY." @@ -5117,6 +5112,11 @@ "message" : [ "TRANSFORM with SERDE is only supported in hive mode." ] + }, + "TRIM_COLLATION" : { + "message" : [ + "TRIM specifier in the collation." + ] } }, "sqlState" : "0A000" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala index 0cff70436db7d..b67e66323bbbd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collationExpressions.scala @@ -52,6 +52,10 @@ object CollateExpressionBuilder extends ExpressionBuilder { if (evalCollation == null) { throw QueryCompilationErrors.unexpectedNullError("collation", collationExpr) } else { + if (!SQLConf.get.trimCollationEnabled && + evalCollation.toString.toUpperCase().contains("TRIM")) { + throw QueryCompilationErrors.trimCollationNotEnabledError() + } Collate(e, evalCollation.toString) } case (_: StringType, false) => throw QueryCompilationErrors.nonFoldableArgumentError( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index 674005caaf1b2..ed6cf329eeca8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -2557,6 +2557,10 @@ class AstBuilder extends DataTypeAstBuilder } override def visitCollateClause(ctx: CollateClauseContext): String = withOrigin(ctx) { + val collationName = ctx.collationName.getText + if (!SQLConf.get.trimCollationEnabled && collationName.toUpperCase().contains("TRIM")) { + throw QueryCompilationErrors.trimCollationNotEnabledError() + } ctx.identifier.getText } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index 0b5255e95f073..0d27f7bedbd3e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -351,6 +351,13 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat ) } + def trimCollationNotEnabledError(): Throwable = { + new AnalysisException( + errorClass = "UNSUPPORTED_FEATURE.TRIM_COLLATION", + messageParameters = Map.empty + ) + } + def unresolvedUsingColForJoinError( colName: String, suggestion: String, side: String): Throwable = { new AnalysisException( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 9c46dd8e83ab2..ea187c0316c17 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -759,6 +759,18 @@ object SQLConf { .checkValue(_ > 0, "The initial number of partitions must be positive.") .createOptional + lazy val TRIM_COLLATION_ENABLED = + buildConf("spark.sql.collation.trim.enabled") + .internal() + .doc( + "Trim collation feature is under development and its use should be done under this" + + "feature flag. Trim collation trims leading, trailing or both spaces depending of" + + "specifier (LTRIM, RTRIM, TRIM)." + ) + .version("4.0.0") + .booleanConf + .createWithDefault(Utils.isTesting) + val DEFAULT_COLLATION = buildConf(SqlApiConfHelper.DEFAULT_COLLATION) .doc("Sets default collation to use for string literals, parameter markers or the string" + @@ -5482,6 +5494,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { } } + def trimCollationEnabled: Boolean = getConf(TRIM_COLLATION_ENABLED) + override def defaultStringType: StringType = { if (getConf(DEFAULT_COLLATION).toUpperCase(Locale.ROOT) == "UTF8_BINARY") { StringType diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala index 1c735154f25ed..8fc860c503c96 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala @@ -168,6 +168,10 @@ class SparkSqlAstBuilder extends AstBuilder { * }}} */ override def visitSetCollation(ctx: SetCollationContext): LogicalPlan = withOrigin(ctx) { + val collationName = ctx.collationName.getText + if (!SQLConf.get.trimCollationEnabled && collationName.toUpperCase().contains("TRIM")) { + throw QueryCompilationErrors.trimCollationNotEnabledError() + } val key = SQLConf.DEFAULT_COLLATION.key SetCommand(Some(key -> Some(ctx.identifier.getText.toUpperCase(Locale.ROOT)))) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala index 632b9305feb57..03d3ed6ac7cb5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala @@ -44,27 +44,57 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { private val allFileBasedDataSources = collationPreservingSources ++ collationNonPreservingSources test("collate returns proper type") { - Seq("utf8_binary", "utf8_lcase", "unicode", "unicode_ci").foreach { collationName => + Seq( + "utf8_binary", + "utf8_lcase", + "unicode", + "unicode_ci", + "unicode_ltrim_ci", + "utf8_lcase_trim", + "utf8_binary_rtrim" + ).foreach { collationName => checkAnswer(sql(s"select 'aaa' collate $collationName"), Row("aaa")) val collationId = CollationFactory.collationNameToId(collationName) - assert(sql(s"select 'aaa' collate $collationName").schema(0).dataType - == StringType(collationId)) + assert( + sql(s"select 'aaa' collate $collationName").schema(0).dataType + == StringType(collationId) + ) } } test("collation name is case insensitive") { - Seq("uTf8_BiNaRy", "utf8_lcase", "uNicOde", "UNICODE_ci").foreach { collationName => + Seq( + "uTf8_BiNaRy", + "utf8_lcase", + "uNicOde", + "UNICODE_ci", + "uNiCoDE_ltRIm_cI", + "UtF8_lCaSE_tRIM", + "utf8_biNAry_RtRiM" + ).foreach { collationName => checkAnswer(sql(s"select 'aaa' collate $collationName"), Row("aaa")) val collationId = CollationFactory.collationNameToId(collationName) - assert(sql(s"select 'aaa' collate $collationName").schema(0).dataType - == StringType(collationId)) + assert( + sql(s"select 'aaa' collate $collationName").schema(0).dataType + == StringType(collationId) + ) } } test("collation expression returns name of collation") { - Seq("utf8_binary", "utf8_lcase", "unicode", "unicode_ci").foreach { collationName => + Seq( + "utf8_binary", + "utf8_lcase", + "unicode", + "unicode_ci", + "unicode_ci_ltrim", + "utf8_lcase_trim", + "utf8_binary_rtrim" + ).foreach { collationName => checkAnswer( - sql(s"select collation('aaa' collate $collationName)"), Row(collationName.toUpperCase())) + sql(s"select collation('aaa' collate $collationName)"), + Row(collationName.toUpperCase()) + ) } } @@ -77,9 +107,15 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { test("collate function syntax with default collation set") { withSQLConf(SqlApiConf.DEFAULT_COLLATION -> "UTF8_LCASE") { - assert(sql(s"select collate('aaa', 'utf8_lcase')").schema(0).dataType == - StringType("UTF8_LCASE")) + assert( + sql(s"select collate('aaa', 'utf8_lcase')").schema(0).dataType == + StringType("UTF8_LCASE") + ) assert(sql(s"select collate('aaa', 'UNICODE')").schema(0).dataType == StringType("UNICODE")) + assert( + sql(s"select collate('aaa', 'UNICODE_TRIM')").schema(0).dataType == + StringType("UNICODE_TRIM") + ) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala index 832e1873af6a4..5abdca326f2fd 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala @@ -868,6 +868,39 @@ class QueryCompilationErrorsSuite "inputTypes" -> "[\"INT\", \"STRING\", \"STRING\"]")) } + test("SPARK-49666: the trim collation feature is off without collate builder call") { + withSQLConf(SQLConf.TRIM_COLLATION_ENABLED.key -> "false") { + Seq( + "CREATE TABLE t(col STRING COLLATE EN_TRIM_CI) USING parquet", + "CREATE TABLE t(col STRING COLLATE UTF8_LCASE_TRIM) USING parquet", + "SELECT 'aaa' COLLATE UNICODE_LTRIM_CI" + ).foreach { sqlText => + checkError( + exception = intercept[AnalysisException](sql(sqlText)), + condition = "UNSUPPORTED_FEATURE.TRIM_COLLATION" + ) + } + } + } + + test("SPARK-49666: the trim collation feature is off with collate builder call") { + withSQLConf(SQLConf.TRIM_COLLATION_ENABLED.key -> "false") { + Seq( + "SELECT collate('aaa', 'UNICODE_TRIM')", + "SELECT collate('aaa', 'UTF8_BINARY_TRIM')", + "SELECT collate('aaa', 'EN_AI_RTRIM')" + ).foreach { sqlText => + checkError( + exception = intercept[AnalysisException](sql(sqlText)), + condition = "UNSUPPORTED_FEATURE.TRIM_COLLATION", + parameters = Map.empty, + context = + ExpectedContext(fragment = sqlText.substring(7), start = 7, stop = sqlText.length - 1) + ) + } + } + } + test("UNSUPPORTED_CALL: call the unsupported method update()") { checkError( exception = intercept[SparkUnsupportedOperationException] { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala index 82795e551b6bf..094c65c63bfdc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala @@ -517,6 +517,13 @@ class SQLConfSuite extends QueryTest with SharedSparkSession { "confName" -> "spark.sql.session.collation.default", "proposals" -> "UNICODE" )) + + withSQLConf(SQLConf.TRIM_COLLATION_ENABLED.key -> "false") { + checkError( + exception = intercept[AnalysisException](sql(s"SET COLLATION UNICODE_CI_TRIM")), + condition = "UNSUPPORTED_FEATURE.TRIM_COLLATION" + ) + } } test("SPARK-43028: config not found error") { From 97ae372634b119b2b67304df67463b95b20febd9 Mon Sep 17 00:00:00 2001 From: Nick Young Date: Mon, 30 Sep 2024 20:44:51 +0800 Subject: [PATCH 175/189] [SPARK-49819] Disable CollapseProject for correlated subqueries in projection over aggregate correctly ### What changes were proposed in this pull request? CollapseProject should block collapsing with an aggregate if any correlated subquery is present. There are other correlated subqueries that are not ScalarSubquery that are not accounted for here. ### Why are the changes needed? Availability issue. ### Does this PR introduce _any_ user-facing change? Previously failing queries will not fail anymore. ### How was this patch tested? UT. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48286 from n-young-db/n-young-db/collapse-project-correlated-subquery-check. Lead-authored-by: Nick Young Co-authored-by: Wenchen Fan Signed-off-by: Wenchen Fan --- .../sql/catalyst/optimizer/Optimizer.scala | 8 +++----- .../org/apache/spark/sql/SubquerySuite.scala | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 7fc12f7d1fc16..fb234c7bda4c3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.catalog.{InMemoryCatalog, SessionCatalog} import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.SubqueryExpression.hasCorrelatedSubquery import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.plans._ @@ -1232,11 +1233,8 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper { * in aggregate if they are also part of the grouping expressions. Otherwise the plan * after subquery rewrite will not be valid. */ - private def canCollapseAggregate(p: Project, a: Aggregate): Boolean = { - p.projectList.forall(_.collect { - case s: ScalarSubquery if s.outerAttrs.nonEmpty => s - }.isEmpty) - } + private def canCollapseAggregate(p: Project, a: Aggregate): Boolean = + !p.projectList.exists(hasCorrelatedSubquery) def buildCleanedProjectList( upper: Seq[NamedExpression], diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index 6e160b4407ca8..f17cf25565145 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -2142,6 +2142,24 @@ class SubquerySuite extends QueryTest } } + test("SPARK-49819: Do not collapse projects with exist subqueries") { + withTempView("v") { + Seq((0, 1), (1, 2)).toDF("c1", "c2").createOrReplaceTempView("v") + checkAnswer( + sql(""" + |SELECT m, CASE WHEN EXISTS (SELECT SUM(c2) FROM v WHERE c1 = m) THEN 1 ELSE 0 END + |FROM (SELECT MIN(c2) AS m FROM v) + |""".stripMargin), + Row(1, 1) :: Nil) + checkAnswer( + sql(""" + |SELECT c, CASE WHEN EXISTS (SELECT SUM(c2) FROM v WHERE c1 = c) THEN 1 ELSE 0 END + |FROM (SELECT c1 AS c FROM v GROUP BY c1) + |""".stripMargin), + Row(0, 1) :: Row(1, 1) :: Nil) + } + } + test("SPARK-37199: deterministic in QueryPlan considers subquery") { val deterministicQueryPlan = sql("select (select 1 as b) as b") .queryExecution.executedPlan From dbfa909422ad82b0428b258671813510caa6eeac Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Mon, 30 Sep 2024 21:15:42 +0800 Subject: [PATCH 176/189] [SPARK-49816][SQL] Should only update out-going-ref-count for referenced outer CTE relation ### What changes were proposed in this pull request? This PR fixes a long-standing reference counting bug in the rule `InlineCTE`. Let's look at the minimal repro: ``` sql( """ |WITH |t1 AS (SELECT 1 col), |t2 AS (SELECT * FROM t1) |SELECT * FROM t2 |""".stripMargin).createTempView("v") // r1 is un-referenced, but it should not decrease the ref count of t2 inside view v. val df = sql( """ |WITH |r1 AS (SELECT * FROM v), |r2 AS (SELECT * FROM v) |SELECT * FROM r2 |""".stripMargin) ``` The logical plan is something like below ``` WithCTE CTEDef r1 View v WithCTE CTEDef t1 OneRowRelation CTEDef t2 CTERef t1 CTERef t2 // main query of the inner WithCTE CTEDef r2 View v // exactly the same as the view v above WithCTE CTEDef t1 OneRowRelation CTEDef t2 CTERef t1 CTERef t2 CTERef r2 // main query of the outer WithCTE ``` Ideally, the ref count of `t1`, `t2` and `r2` should be all `1`. They will be inlined and the final plan is the `OneRowRelation`. However, in `InlineCTE#buildCTEMap`, when we traverse into `CTEDef r1` and hit `CTERef t2`, we mistakenly update the out-going-ref-count of `r1`, which means that `r1` references `t2` and this is totally wrong. Later on, in `InlineCTE#cleanCTEMap`, we find that `r1` is not referenced at all, so we decrease the ref count of its out-going-ref, which is `t2`, and the ref count of `t2` becomes `0`. Finally, in `InlineCTE#inlineCTE`, we leave the plan of `t2` unchanged because its ref count is `0`, and the plan of `t2` contains `CTERef t1`. `t2` is still inlined so we end up with `CTERef t1` as the final plan without the `WithCTE` node. ### Why are the changes needed? bug fix ### Does this PR introduce _any_ user-facing change? Yes, the query failed before and now can work ### How was this patch tested? new test ### Was this patch authored or co-authored using generative AI tooling? no Closes #48284 from cloud-fan/cte. Authored-by: Wenchen Fan Signed-off-by: Wenchen Fan --- .../sql/catalyst/optimizer/InlineCTE.scala | 31 ++++++++++++------- .../org/apache/spark/sql/CTEInlineSuite.scala | 21 +++++++++++++ 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala index 19aa1d96ccd3f..b3384c4e29566 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala @@ -71,13 +71,13 @@ case class InlineCTE( * @param plan The plan to collect the CTEs from * @param cteMap A mutable map that accumulates the CTEs and their reference information by CTE * ids. - * @param outerCTEId While collecting the map we use this optional CTE id to identify the - * current outer CTE. + * @param collectCTERefs A function to collect CTE references so that the caller side can do some + * bookkeeping work. */ private def buildCTEMap( plan: LogicalPlan, cteMap: mutable.Map[Long, CTEReferenceInfo], - outerCTEId: Option[Long] = None): Unit = { + collectCTERefs: CTERelationRef => Unit = _ => ()): Unit = { plan match { case WithCTE(child, cteDefs) => cteDefs.foreach { cteDef => @@ -89,26 +89,35 @@ case class InlineCTE( ) } cteDefs.foreach { cteDef => - buildCTEMap(cteDef, cteMap, Some(cteDef.id)) + buildCTEMap(cteDef, cteMap, ref => { + // A CTE relation can references CTE relations defined before it in the same `WithCTE`. + // Here we update the out-going-ref-count for it, in case this CTE relation is not + // referenced at all and can be optimized out, and we need to decrease the ref counts + // for CTE relations that are referenced by it. + if (cteDefs.exists(_.id == ref.cteId)) { + cteMap(cteDef.id).increaseOutgoingRefCount(ref.cteId, 1) + } + // Similarly, a CTE relation can reference CTE relations defined in the outer `WithCTE`. + // Here we call the `collectCTERefs` function so that the outer CTE can also update the + // out-going-ref-count if needed. + collectCTERefs(ref) + }) } - buildCTEMap(child, cteMap, outerCTEId) + buildCTEMap(child, cteMap, collectCTERefs) case ref: CTERelationRef => cteMap(ref.cteId) = cteMap(ref.cteId).withRefCountIncreased(1) - outerCTEId.foreach { cteId => - cteMap(cteId).increaseOutgoingRefCount(ref.cteId, 1) - } - + collectCTERefs(ref) case _ => if (plan.containsPattern(CTE)) { plan.children.foreach { child => - buildCTEMap(child, cteMap, outerCTEId) + buildCTEMap(child, cteMap, collectCTERefs) } plan.expressions.foreach { expr => if (expr.containsAllPatterns(PLAN_EXPRESSION, CTE)) { expr.foreach { - case e: SubqueryExpression => buildCTEMap(e.plan, cteMap, outerCTEId) + case e: SubqueryExpression => buildCTEMap(e.plan, cteMap, collectCTERefs) case _ => } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala index 7b608b7438c29..7a2ce1d7836b3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala @@ -714,6 +714,27 @@ abstract class CTEInlineSuiteBase |""".stripMargin) checkAnswer(df, Row(1)) } + + test("SPARK-49816: should only update out-going-ref-count for referenced outer CTE relation") { + withView("v") { + sql( + """ + |WITH + |t1 AS (SELECT 1 col), + |t2 AS (SELECT * FROM t1) + |SELECT * FROM t2 + |""".stripMargin).createTempView("v") + // r1 is un-referenced, but it should not decrease the ref count of t2 inside view v. + val df = sql( + """ + |WITH + |r1 AS (SELECT * FROM v), + |r2 AS (SELECT * FROM v) + |SELECT * FROM r2 + |""".stripMargin) + checkAnswer(df, Row(1)) + } + } } class CTEInlineSuiteAEOff extends CTEInlineSuiteBase with DisableAdaptiveExecutionSuite From 3065dd92ab8f36b019c7be06da59d47c1865fe60 Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Mon, 30 Sep 2024 21:31:15 +0800 Subject: [PATCH 177/189] [SPARK-49561][SQL] Add SQL pipe syntax for the PIVOT and UNPIVOT operators ### What changes were proposed in this pull request? This PR adds SQL pipe syntax support for the PIVOT and UNPIVOT operators. For example: ``` CREATE TEMPORARY VIEW courseSales AS SELECT * FROM VALUES ("dotNET", 2012, 10000), ("Java", 2012, 20000), ("dotNET", 2012, 5000), ("dotNET", 2013, 48000), ("Java", 2013, 30000) as courseSales(course, year, earnings); TABLE courseSales |> SELECT `year`, course, earnings |> PIVOT ( SUM(earnings) FOR course IN ('dotNET', 'Java') ); 2012 15000 20000 2013 48000 30000 ``` ### Why are the changes needed? The SQL pipe operator syntax will let users compose queries in a more flexible fashion. ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds a few unit test cases, but mostly relies on golden file test coverage. I did this to make sure the answers are correct as this feature is implemented and also so we can look at the analyzer output plans to ensure they look right as well. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48093 from dtenedor/pipe-pivot. Authored-by: Daniel Tenedorio Signed-off-by: Wenchen Fan --- .../sql/catalyst/parser/SqlBaseParser.g4 | 5 + .../sql/catalyst/parser/AstBuilder.scala | 12 +- .../analyzer-results/pipe-operators.sql.out | 352 ++++++++++++++++++ .../sql-tests/inputs/pipe-operators.sql | 141 +++++++ .../sql-tests/results/pipe-operators.sql.out | 309 +++++++++++++++ .../sql/execution/SparkSqlParserSuite.scala | 39 +- 6 files changed, 849 insertions(+), 9 deletions(-) diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 866634b041280..33ac3249eb663 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -1499,6 +1499,11 @@ version operatorPipeRightSide : selectClause | whereClause + // The following two cases match the PIVOT or UNPIVOT clause, respectively. + // For each one, we add the other clause as an option in order to return high-quality error + // messages in the event that both are present (this is not allowed). + | pivotClause unpivotClause? + | unpivotClause pivotClause? ; // When `SQL_standard_keyword_behavior=true`, there are 2 kinds of keywords in Spark SQL. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index ed6cf329eeca8..e2350474a8708 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -5893,7 +5893,17 @@ class AstBuilder extends DataTypeAstBuilder SubqueryAlias(SubqueryAlias.generateSubqueryName(), left) } withWhereClause(c, withSubqueryAlias) - }.get) + }.getOrElse(Option(ctx.pivotClause()).map { c => + if (ctx.unpivotClause() != null) { + throw QueryParsingErrors.unpivotWithPivotInFromClauseNotAllowedError(ctx) + } + withPivot(c, left) + }.getOrElse(Option(ctx.unpivotClause()).map { c => + if (ctx.pivotClause() != null) { + throw QueryParsingErrors.unpivotWithPivotInFromClauseNotAllowedError(ctx) + } + withUnpivot(c, left) + }.get))) } /** diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out index c44ce153a2f41..8cd062aeb01a3 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out @@ -62,6 +62,74 @@ InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_d +- LocalRelation [col1#x, col2#x] +-- !query +create temporary view courseSales as select * from values + ("dotNET", 2012, 10000), + ("Java", 2012, 20000), + ("dotNET", 2012, 5000), + ("dotNET", 2013, 48000), + ("Java", 2013, 30000) + as courseSales(course, year, earnings) +-- !query analysis +CreateViewCommand `courseSales`, select * from values + ("dotNET", 2012, 10000), + ("Java", 2012, 20000), + ("dotNET", 2012, 5000), + ("dotNET", 2013, 48000), + ("Java", 2013, 30000) + as courseSales(course, year, earnings), false, false, LocalTempView, UNSUPPORTED, true + +- Project [course#x, year#x, earnings#x] + +- SubqueryAlias courseSales + +- LocalRelation [course#x, year#x, earnings#x] + + +-- !query +create temporary view courseEarnings as select * from values + ("dotNET", 15000, 48000, 22500), + ("Java", 20000, 30000, NULL) + as courseEarnings(course, `2012`, `2013`, `2014`) +-- !query analysis +CreateViewCommand `courseEarnings`, select * from values + ("dotNET", 15000, 48000, 22500), + ("Java", 20000, 30000, NULL) + as courseEarnings(course, `2012`, `2013`, `2014`), false, false, LocalTempView, UNSUPPORTED, true + +- Project [course#x, 2012#x, 2013#x, 2014#x] + +- SubqueryAlias courseEarnings + +- LocalRelation [course#x, 2012#x, 2013#x, 2014#x] + + +-- !query +create temporary view courseEarningsAndSales as select * from values + ("dotNET", 15000, NULL, 48000, 1, 22500, 1), + ("Java", 20000, 1, 30000, 2, NULL, NULL) + as courseEarningsAndSales( + course, earnings2012, sales2012, earnings2013, sales2013, earnings2014, sales2014) +-- !query analysis +CreateViewCommand `courseEarningsAndSales`, select * from values + ("dotNET", 15000, NULL, 48000, 1, 22500, 1), + ("Java", 20000, 1, 30000, 2, NULL, NULL) + as courseEarningsAndSales( + course, earnings2012, sales2012, earnings2013, sales2013, earnings2014, sales2014), false, false, LocalTempView, UNSUPPORTED, true + +- Project [course#x, earnings2012#x, sales2012#x, earnings2013#x, sales2013#x, earnings2014#x, sales2014#x] + +- SubqueryAlias courseEarningsAndSales + +- LocalRelation [course#x, earnings2012#x, sales2012#x, earnings2013#x, sales2013#x, earnings2014#x, sales2014#x] + + +-- !query +create temporary view yearsWithComplexTypes as select * from values + (2012, array(1, 1), map('1', 1), struct(1, 'a')), + (2013, array(2, 2), map('2', 2), struct(2, 'b')) + as yearsWithComplexTypes(y, a, m, s) +-- !query analysis +CreateViewCommand `yearsWithComplexTypes`, select * from values + (2012, array(1, 1), map('1', 1), struct(1, 'a')), + (2013, array(2, 2), map('2', 2), struct(2, 'b')) + as yearsWithComplexTypes(y, a, m, s), false, false, LocalTempView, UNSUPPORTED, true + +- Project [y#x, a#x, m#x, s#x] + +- SubqueryAlias yearsWithComplexTypes + +- LocalRelation [y#x, a#x, m#x, s#x] + + -- !query table t |> select 1 as x @@ -569,6 +637,290 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +table courseSales +|> select `year`, course, earnings +|> pivot ( + sum(earnings) + for course in ('dotNET', 'Java') + ) +-- !query analysis +Project [year#x, __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x[0] AS dotNET#xL, __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x[1] AS Java#xL] ++- Aggregate [year#x], [year#x, pivotfirst(course#x, sum(coursesales.earnings)#xL, dotNET, Java, 0, 0) AS __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x] + +- Aggregate [year#x, course#x], [year#x, course#x, sum(earnings#x) AS sum(coursesales.earnings)#xL] + +- Project [year#x, course#x, earnings#x] + +- SubqueryAlias coursesales + +- View (`courseSales`, [course#x, year#x, earnings#x]) + +- Project [cast(course#x as string) AS course#x, cast(year#x as int) AS year#x, cast(earnings#x as int) AS earnings#x] + +- Project [course#x, year#x, earnings#x] + +- SubqueryAlias courseSales + +- LocalRelation [course#x, year#x, earnings#x] + + +-- !query +table courseSales +|> select `year` as y, course as c, earnings as e +|> pivot ( + sum(e) as s, avg(e) as a + for y in (2012 as firstYear, 2013 as secondYear) + ) +-- !query analysis +Project [c#x, __pivot_sum(e) AS s AS `sum(e) AS s`#x[0] AS firstYear_s#xL, __pivot_avg(e) AS a AS `avg(e) AS a`#x[0] AS firstYear_a#x, __pivot_sum(e) AS s AS `sum(e) AS s`#x[1] AS secondYear_s#xL, __pivot_avg(e) AS a AS `avg(e) AS a`#x[1] AS secondYear_a#x] ++- Aggregate [c#x], [c#x, pivotfirst(y#x, sum(e) AS s#xL, 2012, 2013, 0, 0) AS __pivot_sum(e) AS s AS `sum(e) AS s`#x, pivotfirst(y#x, avg(e) AS a#x, 2012, 2013, 0, 0) AS __pivot_avg(e) AS a AS `avg(e) AS a`#x] + +- Aggregate [c#x, y#x], [c#x, y#x, sum(e#x) AS sum(e) AS s#xL, avg(e#x) AS avg(e) AS a#x] + +- Project [pipeselect(year#x) AS y#x, pipeselect(course#x) AS c#x, pipeselect(earnings#x) AS e#x] + +- SubqueryAlias coursesales + +- View (`courseSales`, [course#x, year#x, earnings#x]) + +- Project [cast(course#x as string) AS course#x, cast(year#x as int) AS year#x, cast(earnings#x as int) AS earnings#x] + +- Project [course#x, year#x, earnings#x] + +- SubqueryAlias courseSales + +- LocalRelation [course#x, year#x, earnings#x] + + +-- !query +select course, `year`, y, a +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + max(a) + for (y, course) in ((2012, 'dotNET'), (2013, 'Java')) + ) +-- !query analysis +Aggregate [year#x], [year#x, max(if ((named_struct(y, y#x, course, course#x) <=> cast(named_struct(col1, 2012, col2, dotNET) as struct))) a#x else cast(null as array)) AS {2012, dotNET}#x, max(if ((named_struct(y, y#x, course, course#x) <=> cast(named_struct(col1, 2013, col2, Java) as struct))) a#x else cast(null as array)) AS {2013, Java}#x] ++- Project [course#x, year#x, y#x, a#x] + +- Join Inner, (year#x = y#x) + :- SubqueryAlias coursesales + : +- View (`courseSales`, [course#x, year#x, earnings#x]) + : +- Project [cast(course#x as string) AS course#x, cast(year#x as int) AS year#x, cast(earnings#x as int) AS earnings#x] + : +- Project [course#x, year#x, earnings#x] + : +- SubqueryAlias courseSales + : +- LocalRelation [course#x, year#x, earnings#x] + +- SubqueryAlias yearswithcomplextypes + +- View (`yearsWithComplexTypes`, [y#x, a#x, m#x, s#x]) + +- Project [cast(y#x as int) AS y#x, cast(a#x as array) AS a#x, cast(m#x as map) AS m#x, cast(s#x as struct) AS s#x] + +- Project [y#x, a#x, m#x, s#x] + +- SubqueryAlias yearsWithComplexTypes + +- LocalRelation [y#x, a#x, m#x, s#x] + + +-- !query +select earnings, `year`, s +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + sum(earnings) + for s in ((1, 'a'), (2, 'b')) + ) +-- !query analysis +Project [year#x, __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x[0] AS {1, a}#xL, __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x[1] AS {2, b}#xL] ++- Aggregate [year#x], [year#x, pivotfirst(s#x, sum(coursesales.earnings)#xL, [1,a], [2,b], 0, 0) AS __pivot_sum(coursesales.earnings) AS `sum(coursesales.earnings)`#x] + +- Aggregate [year#x, s#x], [year#x, s#x, sum(earnings#x) AS sum(coursesales.earnings)#xL] + +- Project [earnings#x, year#x, s#x] + +- Join Inner, (year#x = y#x) + :- SubqueryAlias coursesales + : +- View (`courseSales`, [course#x, year#x, earnings#x]) + : +- Project [cast(course#x as string) AS course#x, cast(year#x as int) AS year#x, cast(earnings#x as int) AS earnings#x] + : +- Project [course#x, year#x, earnings#x] + : +- SubqueryAlias courseSales + : +- LocalRelation [course#x, year#x, earnings#x] + +- SubqueryAlias yearswithcomplextypes + +- View (`yearsWithComplexTypes`, [y#x, a#x, m#x, s#x]) + +- Project [cast(y#x as int) AS y#x, cast(a#x as array) AS a#x, cast(m#x as map) AS m#x, cast(s#x as struct) AS s#x] + +- Project [y#x, a#x, m#x, s#x] + +- SubqueryAlias yearsWithComplexTypes + +- LocalRelation [y#x, a#x, m#x, s#x] + + +-- !query +table courseEarnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query analysis +Filter isnotnull(coalesce(earningsYear#x)) ++- Expand [[course#x, 2012, 2012#x], [course#x, 2013, 2013#x], [course#x, 2014, 2014#x]], [course#x, year#x, earningsYear#x] + +- SubqueryAlias courseearnings + +- View (`courseEarnings`, [course#x, 2012#x, 2013#x, 2014#x]) + +- Project [cast(course#x as string) AS course#x, cast(2012#x as int) AS 2012#x, cast(2013#x as int) AS 2013#x, cast(2014#x as int) AS 2014#x] + +- Project [course#x, 2012#x, 2013#x, 2014#x] + +- SubqueryAlias courseEarnings + +- LocalRelation [course#x, 2012#x, 2013#x, 2014#x] + + +-- !query +table courseEarnings +|> unpivot include nulls ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query analysis +Expand [[course#x, 2012, 2012#x], [course#x, 2013, 2013#x], [course#x, 2014, 2014#x]], [course#x, year#x, earningsYear#x] ++- SubqueryAlias courseearnings + +- View (`courseEarnings`, [course#x, 2012#x, 2013#x, 2014#x]) + +- Project [cast(course#x as string) AS course#x, cast(2012#x as int) AS 2012#x, cast(2013#x as int) AS 2013#x, cast(2014#x as int) AS 2014#x] + +- Project [course#x, 2012#x, 2013#x, 2014#x] + +- SubqueryAlias courseEarnings + +- LocalRelation [course#x, 2012#x, 2013#x, 2014#x] + + +-- !query +table courseEarningsAndSales +|> unpivot include nulls ( + (earnings, sales) for `year` in ( + (earnings2012, sales2012) as `2012`, + (earnings2013, sales2013) as `2013`, + (earnings2014, sales2014) as `2014`) + ) +-- !query analysis +Expand [[course#x, 2012, earnings2012#x, sales2012#x], [course#x, 2013, earnings2013#x, sales2013#x], [course#x, 2014, earnings2014#x, sales2014#x]], [course#x, year#x, earnings#x, sales#x] ++- SubqueryAlias courseearningsandsales + +- View (`courseEarningsAndSales`, [course#x, earnings2012#x, sales2012#x, earnings2013#x, sales2013#x, earnings2014#x, sales2014#x]) + +- Project [cast(course#x as string) AS course#x, cast(earnings2012#x as int) AS earnings2012#x, cast(sales2012#x as int) AS sales2012#x, cast(earnings2013#x as int) AS earnings2013#x, cast(sales2013#x as int) AS sales2013#x, cast(earnings2014#x as int) AS earnings2014#x, cast(sales2014#x as int) AS sales2014#x] + +- Project [course#x, earnings2012#x, sales2012#x, earnings2013#x, sales2013#x, earnings2014#x, sales2014#x] + +- SubqueryAlias courseEarningsAndSales + +- LocalRelation [course#x, earnings2012#x, sales2012#x, earnings2013#x, sales2013#x, earnings2014#x, sales2014#x] + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`year`", + "proposal" : "`course`, `earnings`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 49, + "stopIndex" : 111, + "fragment" : "pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )" + } ] +} + + +-- !query +table courseSales +|> pivot ( + sum(earnings) + for `year` in (course, 2013) + ) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NON_LITERAL_PIVOT_VALUES", + "sqlState" : "42K08", + "messageParameters" : { + "expression" : "\"course\"" + } +} + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "NOT_ALLOWED_IN_FROM.UNPIVOT_WITH_PIVOT", + "sqlState" : "42601", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 186, + "fragment" : "table courseSales\n|> select course, earnings\n|> pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )\n unpivot (\n earningsYear for `year` in (`2012`, `2013`, `2014`)\n )" + } ] +} + + +-- !query +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "NOT_ALLOWED_IN_FROM.UNPIVOT_WITH_PIVOT", + "sqlState" : "42601", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 186, + "fragment" : "table courseSales\n|> select course, earnings\n|> unpivot (\n earningsYear for `year` in (`2012`, `2013`, `2014`)\n )\n pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )" + } ] +} + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "PARSE_SYNTAX_ERROR", + "sqlState" : "42601", + "messageParameters" : { + "error" : "'pivot'", + "hint" : "" + } +} + + +-- !query +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "PARSE_SYNTAX_ERROR", + "sqlState" : "42601", + "messageParameters" : { + "error" : "'unpivot'", + "hint" : "" + } +} + + -- !query drop table t -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql index 49a72137ee047..3aa01d472e83f 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql @@ -12,6 +12,30 @@ drop table if exists st; create table st(x int, col struct) using parquet; insert into st values (1, (2, 3)); +create temporary view courseSales as select * from values + ("dotNET", 2012, 10000), + ("Java", 2012, 20000), + ("dotNET", 2012, 5000), + ("dotNET", 2013, 48000), + ("Java", 2013, 30000) + as courseSales(course, year, earnings); + +create temporary view courseEarnings as select * from values + ("dotNET", 15000, 48000, 22500), + ("Java", 20000, 30000, NULL) + as courseEarnings(course, `2012`, `2013`, `2014`); + +create temporary view courseEarningsAndSales as select * from values + ("dotNET", 15000, NULL, 48000, 1, 22500, 1), + ("Java", 20000, 1, 30000, 2, NULL, NULL) + as courseEarningsAndSales( + course, earnings2012, sales2012, earnings2013, sales2013, earnings2014, sales2014); + +create temporary view yearsWithComplexTypes as select * from values + (2012, array(1, 1), map('1', 1), struct(1, 'a')), + (2013, array(2, 2), map('2', 2), struct(2, 'b')) + as yearsWithComplexTypes(y, a, m, s); + -- SELECT operators: positive tests. --------------------------------------- @@ -185,6 +209,123 @@ table t (select x, sum(length(y)) as sum_len from t group by x) |> where sum(length(y)) = 3; +-- Pivot and unpivot operators: positive tests. +----------------------------------------------- + +table courseSales +|> select `year`, course, earnings +|> pivot ( + sum(earnings) + for course in ('dotNET', 'Java') + ); + +table courseSales +|> select `year` as y, course as c, earnings as e +|> pivot ( + sum(e) as s, avg(e) as a + for y in (2012 as firstYear, 2013 as secondYear) + ); + +-- Pivot on multiple pivot columns with aggregate columns of complex data types. +select course, `year`, y, a +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + max(a) + for (y, course) in ((2012, 'dotNET'), (2013, 'Java')) + ); + +-- Pivot on pivot column of struct type. +select earnings, `year`, s +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + sum(earnings) + for s in ((1, 'a'), (2, 'b')) + ); + +table courseEarnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ); + +table courseEarnings +|> unpivot include nulls ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ); + +table courseEarningsAndSales +|> unpivot include nulls ( + (earnings, sales) for `year` in ( + (earnings2012, sales2012) as `2012`, + (earnings2013, sales2013) as `2013`, + (earnings2014, sales2014) as `2014`) + ); + +-- Pivot and unpivot operators: negative tests. +----------------------------------------------- + +-- The PIVOT operator refers to a column 'year' is not available in the input relation. +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ); + +-- Non-literal PIVOT values are not supported. +table courseSales +|> pivot ( + sum(earnings) + for `year` in (course, 2013) + ); + +-- The PIVOT and UNPIVOT clauses are mutually exclusive. +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ); + +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ); + +-- Multiple PIVOT and/or UNPIVOT clauses are not supported in the same pipe operator. +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ); + +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ); + -- Cleanup. ----------- drop table t; diff --git a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out index 38436b0941034..2c6abe2a277ad 100644 --- a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out @@ -71,6 +71,54 @@ struct<> +-- !query +create temporary view courseSales as select * from values + ("dotNET", 2012, 10000), + ("Java", 2012, 20000), + ("dotNET", 2012, 5000), + ("dotNET", 2013, 48000), + ("Java", 2013, 30000) + as courseSales(course, year, earnings) +-- !query schema +struct<> +-- !query output + + + +-- !query +create temporary view courseEarnings as select * from values + ("dotNET", 15000, 48000, 22500), + ("Java", 20000, 30000, NULL) + as courseEarnings(course, `2012`, `2013`, `2014`) +-- !query schema +struct<> +-- !query output + + + +-- !query +create temporary view courseEarningsAndSales as select * from values + ("dotNET", 15000, NULL, 48000, 1, 22500, 1), + ("Java", 20000, 1, 30000, 2, NULL, NULL) + as courseEarningsAndSales( + course, earnings2012, sales2012, earnings2013, sales2013, earnings2014, sales2014) +-- !query schema +struct<> +-- !query output + + + +-- !query +create temporary view yearsWithComplexTypes as select * from values + (2012, array(1, 1), map('1', 1), struct(1, 'a')), + (2013, array(2, 2), map('2', 2), struct(2, 'b')) + as yearsWithComplexTypes(y, a, m, s) +-- !query schema +struct<> +-- !query output + + + -- !query table t |> select 1 as x @@ -552,6 +600,267 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +table courseSales +|> select `year`, course, earnings +|> pivot ( + sum(earnings) + for course in ('dotNET', 'Java') + ) +-- !query schema +struct +-- !query output +2012 15000 20000 +2013 48000 30000 + + +-- !query +table courseSales +|> select `year` as y, course as c, earnings as e +|> pivot ( + sum(e) as s, avg(e) as a + for y in (2012 as firstYear, 2013 as secondYear) + ) +-- !query schema +struct +-- !query output +Java 20000 20000.0 30000 30000.0 +dotNET 15000 7500.0 48000 48000.0 + + +-- !query +select course, `year`, y, a +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + max(a) + for (y, course) in ((2012, 'dotNET'), (2013, 'Java')) + ) +-- !query schema +struct,{2013, Java}:array> +-- !query output +2012 [1,1] NULL +2013 NULL [2,2] + + +-- !query +select earnings, `year`, s +from courseSales +join yearsWithComplexTypes on `year` = y +|> pivot ( + sum(earnings) + for s in ((1, 'a'), (2, 'b')) + ) +-- !query schema +struct +-- !query output +2012 35000 NULL +2013 NULL 78000 + + +-- !query +table courseEarnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query schema +struct +-- !query output +Java 2012 20000 +Java 2013 30000 +dotNET 2012 15000 +dotNET 2013 48000 +dotNET 2014 22500 + + +-- !query +table courseEarnings +|> unpivot include nulls ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query schema +struct +-- !query output +Java 2012 20000 +Java 2013 30000 +Java 2014 NULL +dotNET 2012 15000 +dotNET 2013 48000 +dotNET 2014 22500 + + +-- !query +table courseEarningsAndSales +|> unpivot include nulls ( + (earnings, sales) for `year` in ( + (earnings2012, sales2012) as `2012`, + (earnings2013, sales2013) as `2013`, + (earnings2014, sales2014) as `2014`) + ) +-- !query schema +struct +-- !query output +Java 2012 20000 1 +Java 2013 30000 2 +Java 2014 NULL NULL +dotNET 2012 15000 NULL +dotNET 2013 48000 1 +dotNET 2014 22500 1 + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "sqlState" : "42703", + "messageParameters" : { + "objectName" : "`year`", + "proposal" : "`course`, `earnings`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 49, + "stopIndex" : 111, + "fragment" : "pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )" + } ] +} + + +-- !query +table courseSales +|> pivot ( + sum(earnings) + for `year` in (course, 2013) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NON_LITERAL_PIVOT_VALUES", + "sqlState" : "42K08", + "messageParameters" : { + "expression" : "\"course\"" + } +} + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "NOT_ALLOWED_IN_FROM.UNPIVOT_WITH_PIVOT", + "sqlState" : "42601", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 186, + "fragment" : "table courseSales\n|> select course, earnings\n|> pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )\n unpivot (\n earningsYear for `year` in (`2012`, `2013`, `2014`)\n )" + } ] +} + + +-- !query +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "NOT_ALLOWED_IN_FROM.UNPIVOT_WITH_PIVOT", + "sqlState" : "42601", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 1, + "stopIndex" : 186, + "fragment" : "table courseSales\n|> select course, earnings\n|> unpivot (\n earningsYear for `year` in (`2012`, `2013`, `2014`)\n )\n pivot (\n sum(earnings)\n for `year` in (2012, 2013)\n )" + } ] +} + + +-- !query +table courseSales +|> select course, earnings +|> pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "PARSE_SYNTAX_ERROR", + "sqlState" : "42601", + "messageParameters" : { + "error" : "'pivot'", + "hint" : "" + } +} + + +-- !query +table courseSales +|> select course, earnings +|> unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + unpivot ( + earningsYear for `year` in (`2012`, `2013`, `2014`) + ) + pivot ( + sum(earnings) + for `year` in (2012, 2013) + ) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "PARSE_SYNTAX_ERROR", + "sqlState" : "42601", + "messageParameters" : { + "error" : "'unpivot'", + "hint" : "" + } +} + + -- !query drop table t -- !query schema diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala index ab949c5a21e44..1111a65c6a526 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.analysis.{AnalysisTest, UnresolvedAlias, Un import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Concat, GreaterThan, Literal, NullsFirst, SortOrder, UnresolvedWindowExpression, UnspecifiedFrame, WindowSpecDefinition, WindowSpecReference} import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.catalyst.plans.logical._ -import org.apache.spark.sql.catalyst.trees.TreePattern.{FILTER, LOCAL_RELATION, PROJECT, UNRESOLVED_RELATION} +import org.apache.spark.sql.catalyst.trees.TreePattern._ import org.apache.spark.sql.connector.catalog.TableCatalog import org.apache.spark.sql.execution.command._ import org.apache.spark.sql.execution.datasources.{CreateTempViewUsing, RefreshResource} @@ -887,24 +887,47 @@ class SparkSqlParserSuite extends AnalysisTest with SharedSparkSession { // Basic selection. // Here we check that every parsed plan contains a projection and a source relation or // inline table. - def checkPipeSelect(query: String): Unit = { + def check(query: String, patterns: Seq[TreePattern]): Unit = { val plan: LogicalPlan = parser.parsePlan(query) - assert(plan.containsPattern(PROJECT)) + assert(patterns.exists(plan.containsPattern)) assert(plan.containsAnyPattern(UNRESOLVED_RELATION, LOCAL_RELATION)) } + def checkPipeSelect(query: String): Unit = check(query, Seq(PROJECT)) checkPipeSelect("TABLE t |> SELECT 1 AS X") checkPipeSelect("TABLE t |> SELECT 1 AS X, 2 AS Y |> SELECT X + Y AS Z") checkPipeSelect("VALUES (0), (1) tab(col) |> SELECT col * 2 AS result") // Basic WHERE operators. - def checkPipeWhere(query: String): Unit = { - val plan: LogicalPlan = parser.parsePlan(query) - assert(plan.containsPattern(FILTER)) - assert(plan.containsAnyPattern(UNRESOLVED_RELATION, LOCAL_RELATION)) - } + def checkPipeWhere(query: String): Unit = check(query, Seq(FILTER)) checkPipeWhere("TABLE t |> WHERE X = 1") checkPipeWhere("TABLE t |> SELECT X, LENGTH(Y) AS Z |> WHERE X + LENGTH(Y) < 4") checkPipeWhere("TABLE t |> WHERE X = 1 AND Y = 2 |> WHERE X + Y = 3") checkPipeWhere("VALUES (0), (1) tab(col) |> WHERE col < 1") + // PIVOT and UNPIVOT operations + def checkPivotUnpivot(query: String): Unit = check(query, Seq(PIVOT, UNPIVOT)) + checkPivotUnpivot( + """ + |SELECT * FROM VALUES + | ("dotNET", 2012, 10000), + | ("Java", 2012, 20000), + | ("dotNET", 2012, 5000), + | ("dotNET", 2013, 48000), + | ("Java", 2013, 30000) + | AS courseSales(course, year, earnings) + ||> PIVOT ( + | SUM(earnings) + | FOR course IN ('dotNET', 'Java') + |) + |""".stripMargin) + checkPivotUnpivot( + """ + |SELECT * FROM VALUES + | ("dotNET", 15000, 48000, 22500), + | ("Java", 20000, 30000, NULL) + | AS courseEarnings(course, `2012`, `2013`, `2014`) + ||> UNPIVOT ( + | earningsYear FOR year IN (`2012`, `2013`, `2014`) + |) + |""".stripMargin) } } } From a7fa2700e0f0f70ec6306f48a5bd137225029b80 Mon Sep 17 00:00:00 2001 From: Julek Sompolski Date: Mon, 30 Sep 2024 23:39:50 +0800 Subject: [PATCH 178/189] [SPARK-48196][SQL] Turn QueryExecution lazy val plans into LazyTry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? Currently, when evaluation of `lazy val` of some of the plans fails in QueryExecution, this `lazy val` remains not initialized, and another attempt will be made to initialize it the next time it's referenced. This leads to planning being performed multiple times, resulting in inefficiencies, and potential duplication of side effects, for example from ConvertToLocalRelation that can pull in UDFs with side effects. ### Why are the changes needed? Current behaviour leads to inefficiencies and subtle problems in accidental situations, for example when plans are accessed for logging purposes. ### Does this PR introduce _any_ user-facing change? Yes. This change would bring slight behaviour changes: Examples: ``` val df = a.join(b) spark.conf.set(“spark.sql.crossJoin.enabled”, “false”) try { df.collect() } catch { case _ => } spark.conf.set(“spark.sql.crossJoin.enabled”, “true”) df.collect() ``` This used to succeed, because the first time around the plan will not be initialized because it threw an error because of the cartprod, and the second time around it will try to initialize it again and pick up the new config. This will now fail, because the second execution will retrieve the error from the first time around instead of retrying. The old semantics is if plan evaluation fails, try again next time it's accessed and if plan evaluation ever succeeded, keep that plan. The new semantics is that if plan evaluation fails, it keeps that error and rethrows it next time the plan is accessed. A new QueryExecution object / new Dataset is needed to reset it. Spark 4.0 may be a good candidate for a slight change in this, to make sure that we don't re-execute the optimizer, and potential side effects of it. Note: These behaviour changes have already happened in Spark Connect mode, where the Dataset object is not reused across execution. This change makes Spark Classic and Spark Connect behave the same again. ### How was this patch tested? Existing tests shows no issues, except for the tests that exhibit the behaviour change described above. ### Was this patch authored or co-authored using generative AI tooling? Trivial code completion suggestions. Generated-by: Github Copilot Closes #48211 from juliuszsompolski/SPARK-48196-lazyplans. Authored-by: Julek Sompolski Signed-off-by: Wenchen Fan --- python/pyspark/sql/tests/test_udf.py | 3 +- .../spark/sql/execution/QueryExecution.scala | 63 ++++++++++++------- 2 files changed, 44 insertions(+), 22 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf.py b/python/pyspark/sql/tests/test_udf.py index 6f672b0ae5fb3..879329bd80c0b 100644 --- a/python/pyspark/sql/tests/test_udf.py +++ b/python/pyspark/sql/tests/test_udf.py @@ -237,11 +237,12 @@ def test_udf_in_join_condition(self): f = udf(lambda a, b: a == b, BooleanType()) # The udf uses attributes from both sides of join, so it is pulled out as Filter + # Cross join. - df = left.join(right, f("a", "b")) with self.sql_conf({"spark.sql.crossJoin.enabled": False}): + df = left.join(right, f("a", "b")) with self.assertRaisesRegex(AnalysisException, "Detected implicit cartesian product"): df.collect() with self.sql_conf({"spark.sql.crossJoin.enabled": True}): + df = left.join(right, f("a", "b")) self.assertEqual(df.collect(), [Row(a=1, b=1)]) def test_udf_in_left_outer_join_condition(self): diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 5c894eb7555b1..6ff2c5d4b9d32 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -46,8 +46,8 @@ import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery import org.apache.spark.sql.execution.streaming.{IncrementalExecution, OffsetSeqMetadata, WatermarkPropagator} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.OutputMode +import org.apache.spark.util.{LazyTry, Utils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.Utils /** * The primary workflow for executing relational queries using Spark. Designed to allow easy @@ -86,7 +86,7 @@ class QueryExecution( } } - lazy val analyzed: LogicalPlan = { + private val lazyAnalyzed = LazyTry { val plan = executePhase(QueryPlanningTracker.ANALYSIS) { // We can't clone `logical` here, which will reset the `_analyzed` flag. sparkSession.sessionState.analyzer.executeAndCheck(logical, tracker) @@ -95,12 +95,18 @@ class QueryExecution( plan } - lazy val commandExecuted: LogicalPlan = mode match { - case CommandExecutionMode.NON_ROOT => analyzed.mapChildren(eagerlyExecuteCommands) - case CommandExecutionMode.ALL => eagerlyExecuteCommands(analyzed) - case CommandExecutionMode.SKIP => analyzed + def analyzed: LogicalPlan = lazyAnalyzed.get + + private val lazyCommandExecuted = LazyTry { + mode match { + case CommandExecutionMode.NON_ROOT => analyzed.mapChildren(eagerlyExecuteCommands) + case CommandExecutionMode.ALL => eagerlyExecuteCommands(analyzed) + case CommandExecutionMode.SKIP => analyzed + } } + def commandExecuted: LogicalPlan = lazyCommandExecuted.get + private def commandExecutionName(command: Command): String = command match { case _: CreateTableAsSelect => "create" case _: ReplaceTableAsSelect => "replace" @@ -141,22 +147,28 @@ class QueryExecution( } } - // The plan that has been normalized by custom rules, so that it's more likely to hit cache. - lazy val normalized: LogicalPlan = { + private val lazyNormalized = LazyTry { QueryExecution.normalize(sparkSession, commandExecuted, Some(tracker)) } - lazy val withCachedData: LogicalPlan = sparkSession.withActive { - assertAnalyzed() - assertSupported() - // clone the plan to avoid sharing the plan instance between different stages like analyzing, - // optimizing and planning. - sparkSession.sharedState.cacheManager.useCachedData(normalized.clone()) + // The plan that has been normalized by custom rules, so that it's more likely to hit cache. + def normalized: LogicalPlan = lazyNormalized.get + + private val lazyWithCachedData = LazyTry { + sparkSession.withActive { + assertAnalyzed() + assertSupported() + // clone the plan to avoid sharing the plan instance between different stages like analyzing, + // optimizing and planning. + sparkSession.sharedState.cacheManager.useCachedData(normalized.clone()) + } } + def withCachedData: LogicalPlan = lazyWithCachedData.get + def assertCommandExecuted(): Unit = commandExecuted - lazy val optimizedPlan: LogicalPlan = { + private val lazyOptimizedPlan = LazyTry { // We need to materialize the commandExecuted here because optimizedPlan is also tracked under // the optimizing phase assertCommandExecuted() @@ -174,9 +186,11 @@ class QueryExecution( } } + def optimizedPlan: LogicalPlan = lazyOptimizedPlan.get + def assertOptimized(): Unit = optimizedPlan - lazy val sparkPlan: SparkPlan = { + private val lazySparkPlan = LazyTry { // We need to materialize the optimizedPlan here because sparkPlan is also tracked under // the planning phase assertOptimized() @@ -187,11 +201,11 @@ class QueryExecution( } } + def sparkPlan: SparkPlan = lazySparkPlan.get + def assertSparkPlanPrepared(): Unit = sparkPlan - // executedPlan should not be used to initialize any SparkPlan. It should be - // only used for execution. - lazy val executedPlan: SparkPlan = { + private val lazyExecutedPlan = LazyTry { // We need to materialize the optimizedPlan here, before tracking the planning phase, to ensure // that the optimization time is not counted as part of the planning phase. assertOptimized() @@ -206,8 +220,16 @@ class QueryExecution( plan } + // executedPlan should not be used to initialize any SparkPlan. It should be + // only used for execution. + def executedPlan: SparkPlan = lazyExecutedPlan.get + def assertExecutedPlanPrepared(): Unit = executedPlan + val lazyToRdd = LazyTry { + new SQLExecutionRDD(executedPlan.execute(), sparkSession.sessionState.conf) + } + /** * Internal version of the RDD. Avoids copies and has no schema. * Note for callers: Spark may apply various optimization including reusing object: this means @@ -218,8 +240,7 @@ class QueryExecution( * Given QueryExecution is not a public class, end users are discouraged to use this: please * use `Dataset.rdd` instead where conversion will be applied. */ - lazy val toRdd: RDD[InternalRow] = new SQLExecutionRDD( - executedPlan.execute(), sparkSession.sessionState.conf) + def toRdd: RDD[InternalRow] = lazyToRdd.get /** Get the metrics observed during the execution of the query plan. */ def observedMetrics: Map[String, Row] = CollectMetricsExec.collect(executedPlan) From d68048b06a046cc67ff431fdd8a687b0a1f43603 Mon Sep 17 00:00:00 2001 From: prathit06 Date: Mon, 30 Sep 2024 14:26:01 -0700 Subject: [PATCH 179/189] [SPARK-49833][K8S] Support user-defined annotations for OnDemand PVCs ### What changes were proposed in this pull request? Currently for on-demand PVCs we cannot add user-defined annotations, user-defined annotations can greatly help to add tags in underlying storage. For e.g. if we add `k8s-pvc-tagger/tags` annotation & provide a map like {"env":"dev"}, the same tags are reflected on underlying storage (for e.g. AWS EBS) ### Why are the changes needed? Changes are needed so users can set custom annotations to PVCs ### Does this PR introduce _any_ user-facing change? It does not break any existing behaviour but adds a new feature/improvement to enable custom annotations additions to ondemand PVCs ### How was this patch tested? This was tested in internal/production k8 cluster ### Was this patch authored or co-authored using generative AI tooling? No Closes #48299 from prathit06/ondemand-pvc-annotations. Authored-by: prathit06 Signed-off-by: Dongjoon Hyun --- docs/running-on-kubernetes.md | 18 +++++ .../org/apache/spark/deploy/k8s/Config.scala | 1 + .../deploy/k8s/KubernetesVolumeSpec.scala | 3 +- .../deploy/k8s/KubernetesVolumeUtils.scala | 14 +++- .../features/MountVolumesFeatureStep.scala | 7 +- .../spark/deploy/k8s/KubernetesTestConf.scala | 8 +- .../k8s/KubernetesVolumeUtilsSuite.scala | 42 +++++++++- .../MountVolumesFeatureStepSuite.scala | 77 +++++++++++++++++++ 8 files changed, 160 insertions(+), 10 deletions(-) diff --git a/docs/running-on-kubernetes.md b/docs/running-on-kubernetes.md index d8be32e047717..f8b935fd77f5c 100644 --- a/docs/running-on-kubernetes.md +++ b/docs/running-on-kubernetes.md @@ -1191,6 +1191,15 @@ See the [configuration page](configuration.html) for information on Spark config 4.0.0 + + spark.kubernetes.driver.volumes.[VolumeType].[VolumeName].annotation.[AnnotationName] + (none) + + Configure Kubernetes Volume annotations passed to the Kubernetes with AnnotationName as key having specified value, must conform with Kubernetes annotations format. For example, + spark.kubernetes.driver.volumes.persistentVolumeClaim.checkpointpvc.annotation.foo=bar. + + 4.0.0 + spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].mount.path (none) @@ -1236,6 +1245,15 @@ See the [configuration page](configuration.html) for information on Spark config 4.0.0 + + spark.kubernetes.executor.volumes.[VolumeType].[VolumeName].annotation.[AnnotationName] + (none) + + Configure Kubernetes Volume annotations passed to the Kubernetes with AnnotationName as key having specified value, must conform with Kubernetes annotations format. For example, + spark.kubernetes.executor.volumes.persistentVolumeClaim.checkpointpvc.annotation.foo=bar. + + 4.0.0 + spark.kubernetes.local.dirs.tmpfs false diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala index 9c50f8ddb00cc..db7fc85976c2a 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Config.scala @@ -779,6 +779,7 @@ private[spark] object Config extends Logging { val KUBERNETES_VOLUMES_OPTIONS_SIZE_LIMIT_KEY = "options.sizeLimit" val KUBERNETES_VOLUMES_OPTIONS_SERVER_KEY = "options.server" val KUBERNETES_VOLUMES_LABEL_KEY = "label." + val KUBERNETES_VOLUMES_ANNOTATION_KEY = "annotation." val KUBERNETES_DRIVER_ENV_PREFIX = "spark.kubernetes.driverEnv." val KUBERNETES_DNS_SUBDOMAIN_NAME_MAX_LENGTH = 253 diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala index b4fe414e3cde5..b7113a562fa06 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeSpec.scala @@ -25,7 +25,8 @@ private[spark] case class KubernetesPVCVolumeConf( claimName: String, storageClass: Option[String] = None, size: Option[String] = None, - labels: Option[Map[String, String]] = None) + labels: Option[Map[String, String]] = None, + annotations: Option[Map[String, String]] = None) extends KubernetesVolumeSpecificConf private[spark] case class KubernetesEmptyDirVolumeConf( diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala index 88bb998d88b7d..95821a909f351 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtils.scala @@ -47,6 +47,7 @@ object KubernetesVolumeUtils { val subPathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_SUBPATH_KEY" val subPathExprKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_MOUNT_SUBPATHEXPR_KEY" val labelKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_LABEL_KEY" + val annotationKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_ANNOTATION_KEY" verifyMutuallyExclusiveOptionKeys(properties, subPathKey, subPathExprKey) val volumeLabelsMap = properties @@ -54,6 +55,11 @@ object KubernetesVolumeUtils { .map { case (k, v) => k.replaceAll(labelKey, "") -> v } + val volumeAnnotationsMap = properties + .filter(_._1.startsWith(annotationKey)) + .map { + case (k, v) => k.replaceAll(annotationKey, "") -> v + } KubernetesVolumeSpec( volumeName = volumeName, @@ -62,7 +68,7 @@ object KubernetesVolumeUtils { mountSubPathExpr = properties.getOrElse(subPathExprKey, ""), mountReadOnly = properties.get(readOnlyKey).exists(_.toBoolean), volumeConf = parseVolumeSpecificConf(properties, - volumeType, volumeName, Option(volumeLabelsMap))) + volumeType, volumeName, Option(volumeLabelsMap), Option(volumeAnnotationsMap))) }.toSeq } @@ -86,7 +92,8 @@ object KubernetesVolumeUtils { options: Map[String, String], volumeType: String, volumeName: String, - labels: Option[Map[String, String]]): KubernetesVolumeSpecificConf = { + labels: Option[Map[String, String]], + annotations: Option[Map[String, String]]): KubernetesVolumeSpecificConf = { volumeType match { case KUBERNETES_VOLUMES_HOSTPATH_TYPE => val pathKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_PATH_KEY" @@ -107,7 +114,8 @@ object KubernetesVolumeUtils { options(claimNameKey), options.get(storageClassKey), options.get(sizeLimitKey), - labels) + labels, + annotations) case KUBERNETES_VOLUMES_EMPTYDIR_TYPE => val mediumKey = s"$volumeType.$volumeName.$KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY" diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala index eea4604010b21..3d89696f19fcc 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStep.scala @@ -74,7 +74,7 @@ private[spark] class MountVolumesFeatureStep(conf: KubernetesConf) new VolumeBuilder() .withHostPath(new HostPathVolumeSource(hostPath, volumeType)) - case KubernetesPVCVolumeConf(claimNameTemplate, storageClass, size, labels) => + case KubernetesPVCVolumeConf(claimNameTemplate, storageClass, size, labels, annotations) => val claimName = conf match { case c: KubernetesExecutorConf => claimNameTemplate @@ -91,12 +91,17 @@ private[spark] class MountVolumesFeatureStep(conf: KubernetesConf) case Some(customLabelsMap) => (customLabelsMap ++ defaultVolumeLabels).asJava case None => defaultVolumeLabels.asJava } + val volumeAnnotations = annotations match { + case Some(value) => value.asJava + case None => Map[String, String]().asJava + } additionalResources.append(new PersistentVolumeClaimBuilder() .withKind(PVC) .withApiVersion("v1") .withNewMetadata() .withName(claimName) .addToLabels(volumeLabels) + .addToAnnotations(volumeAnnotations) .endMetadata() .withNewSpec() .withStorageClassName(storageClass.get) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala index e0ddcd3d416f0..e5ed79718d733 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesTestConf.scala @@ -118,7 +118,7 @@ object KubernetesTestConf { KUBERNETES_VOLUMES_OPTIONS_PATH_KEY -> hostPath, KUBERNETES_VOLUMES_OPTIONS_TYPE_KEY -> volumeType)) - case KubernetesPVCVolumeConf(claimName, storageClass, sizeLimit, labels) => + case KubernetesPVCVolumeConf(claimName, storageClass, sizeLimit, labels, annotations) => val sconf = storageClass .map { s => (KUBERNETES_VOLUMES_OPTIONS_CLAIM_STORAGE_CLASS_KEY, s) }.toMap val lconf = sizeLimit.map { l => (KUBERNETES_VOLUMES_OPTIONS_SIZE_LIMIT_KEY, l) }.toMap @@ -126,9 +126,13 @@ object KubernetesTestConf { case Some(value) => value.map { case(k, v) => s"label.$k" -> v } case None => Map() } + val aannotations = annotations match { + case Some(value) => value.map { case (k, v) => s"annotation.$k" -> v } + case None => Map() + } (KUBERNETES_VOLUMES_PVC_TYPE, Map(KUBERNETES_VOLUMES_OPTIONS_CLAIM_NAME_KEY -> claimName) ++ - sconf ++ lconf ++ llabels) + sconf ++ lconf ++ llabels ++ aannotations) case KubernetesEmptyDirVolumeConf(medium, sizeLimit) => val mconf = medium.map { m => (KUBERNETES_VOLUMES_OPTIONS_MEDIUM_KEY, m) }.toMap diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala index 1e62db725fb6e..3c57cba9a7ff0 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/KubernetesVolumeUtilsSuite.scala @@ -96,7 +96,7 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.mountPath === "/path") assert(volumeSpec.mountReadOnly) assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === - KubernetesPVCVolumeConf("claimName", labels = Some(Map()))) + KubernetesPVCVolumeConf("claimName", labels = Some(Map()), annotations = Some(Map()))) } test("SPARK-49598: Parses persistentVolumeClaim volumes correctly with labels") { @@ -113,7 +113,8 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.mountReadOnly) assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === KubernetesPVCVolumeConf(claimName = "claimName", - labels = Some(Map("env" -> "test", "foo" -> "bar")))) + labels = Some(Map("env" -> "test", "foo" -> "bar")), + annotations = Some(Map()))) } test("SPARK-49598: Parses persistentVolumeClaim volumes & puts " + @@ -128,7 +129,8 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { assert(volumeSpec.mountPath === "/path") assert(volumeSpec.mountReadOnly) assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === - KubernetesPVCVolumeConf(claimName = "claimName", labels = Some(Map()))) + KubernetesPVCVolumeConf(claimName = "claimName", labels = Some(Map()), + annotations = Some(Map()))) } test("Parses emptyDir volumes correctly") { @@ -280,4 +282,38 @@ class KubernetesVolumeUtilsSuite extends SparkFunSuite { }.getMessage assert(m.contains("smaller than 1KiB. Missing units?")) } + + test("SPARK-49833: Parses persistentVolumeClaim volumes correctly with annotations") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.path", "/path") + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.readOnly", "true") + sparkConf.set("test.persistentVolumeClaim.volumeName.options.claimName", "claimName") + sparkConf.set("test.persistentVolumeClaim.volumeName.annotation.key1", "value1") + sparkConf.set("test.persistentVolumeClaim.volumeName.annotation.key2", "value2") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.mountReadOnly) + assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === + KubernetesPVCVolumeConf(claimName = "claimName", + labels = Some(Map()), + annotations = Some(Map("key1" -> "value1", "key2" -> "value2")))) + } + + test("SPARK-49833: Parses persistentVolumeClaim volumes & puts " + + "annotations as empty Map if not provided") { + val sparkConf = new SparkConf(false) + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.path", "/path") + sparkConf.set("test.persistentVolumeClaim.volumeName.mount.readOnly", "true") + sparkConf.set("test.persistentVolumeClaim.volumeName.options.claimName", "claimName") + + val volumeSpec = KubernetesVolumeUtils.parseVolumesWithPrefix(sparkConf, "test.").head + assert(volumeSpec.volumeName === "volumeName") + assert(volumeSpec.mountPath === "/path") + assert(volumeSpec.mountReadOnly) + assert(volumeSpec.volumeConf.asInstanceOf[KubernetesPVCVolumeConf] === + KubernetesPVCVolumeConf(claimName = "claimName", labels = Some(Map()), + annotations = Some(Map()))) + } } diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala index c94a7a6ec26a7..293773ddb9ec5 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/MountVolumesFeatureStepSuite.scala @@ -496,4 +496,81 @@ class MountVolumesFeatureStepSuite extends SparkFunSuite { assert(mounts(1).getMountPath === "/tmp/bar") assert(mounts(1).getSubPath === "bar") } + + test("SPARK-49833: Create and mounts persistentVolumeClaims in driver with annotations") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + "", + true, + KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, + storageClass = Some("gp3"), + size = Some("1Mi"), + annotations = Some(Map("env" -> "test"))) + ) + + val kubernetesConf = KubernetesTestConf.createDriverConf(volumes = Seq(volumeConf)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + assert(configuredPod.pod.getSpec.getVolumes.size() === 1) + val pvcClaim = configuredPod.pod.getSpec.getVolumes.get(0).getPersistentVolumeClaim + assert(pvcClaim.getClaimName.endsWith("-driver-pvc-0")) + } + + test("SPARK-49833: Create and mounts persistentVolumeClaims in executors with annotations") { + val volumeConf = KubernetesVolumeSpec( + "testVolume", + "/tmp", + "", + "", + true, + KubernetesPVCVolumeConf(claimName = MountVolumesFeatureStep.PVC_ON_DEMAND, + storageClass = Some("gp3"), + size = Some("1Mi"), + annotations = Some(Map("env" -> "exec-test"))) + ) + + val executorConf = KubernetesTestConf.createExecutorConf(volumes = Seq(volumeConf)) + val executorStep = new MountVolumesFeatureStep(executorConf) + val executorPod = executorStep.configurePod(SparkPod.initialPod()) + + assert(executorPod.pod.getSpec.getVolumes.size() === 1) + val executorPVC = executorPod.pod.getSpec.getVolumes.get(0).getPersistentVolumeClaim + assert(executorPVC.getClaimName.endsWith("-exec-1-pvc-0")) + } + + test("SPARK-49833: Mount multiple volumes to executor with annotations") { + val pvcVolumeConf1 = KubernetesVolumeSpec( + "checkpointVolume1", + "/checkpoints1", + "", + "", + true, + KubernetesPVCVolumeConf(claimName = "pvcClaim1", + storageClass = Some("gp3"), + size = Some("1Mi"), + annotations = Some(Map("env1" -> "exec-test-1"))) + ) + + val pvcVolumeConf2 = KubernetesVolumeSpec( + "checkpointVolume2", + "/checkpoints2", + "", + "", + true, + KubernetesPVCVolumeConf(claimName = "pvcClaim2", + storageClass = Some("gp3"), + size = Some("1Mi"), + annotations = Some(Map("env2" -> "exec-test-2"))) + ) + + val kubernetesConf = KubernetesTestConf.createExecutorConf( + volumes = Seq(pvcVolumeConf1, pvcVolumeConf2)) + val step = new MountVolumesFeatureStep(kubernetesConf) + val configuredPod = step.configurePod(SparkPod.initialPod()) + + assert(configuredPod.pod.getSpec.getVolumes.size() === 2) + assert(configuredPod.container.getVolumeMounts.size() === 2) + } } From 123361137bbe4db4120111777091829c5abc807a Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Mon, 30 Sep 2024 15:56:58 -0700 Subject: [PATCH 180/189] [SPARK-49732][CORE][K8S] Spark deamons should respect `spark.log.structuredLogging.enabled` conf ### What changes were proposed in this pull request? Explicitly call `Logging.uninitialize()` after `SparkConf` loading `spark-defaults.conf` ### Why are the changes needed? SPARK-49015 fixes a similar issue that affects services started through `SparkSubmit`, while for other services like SHS, there is still a chance that the logging system is initialized before `SparkConf` constructed, so `spark.log.structuredLogging.enabled` configured at `spark-defaults.conf` won't take effect. The issue only happens when the logging system is initialized before `SparkConf` loading `spark-defaults.conf`. [example 1](https://github.com/apache/spark/pull/47500#issuecomment-2320426384), when `java.net.InetAddress.getLocalHost` returns `127.0.0.1`, ``` scala> java.net.InetAddress.getLocalHost res0: java.net.InetAddress = H27212-MAC-01.local/127.0.0.1 ``` the logging system will be initialized early. ``` {"ts":"2024-09-22T12:50:37.082Z","level":"WARN","msg":"Your hostname, H27212-MAC-01.local, resolves to a loopback address: 127.0.0.1; using 192.168.32.130 instead (on interface en0)","context":{"host":"H27212-MAC-01.local","host_port":"127.0.0.1","host_port2":"192.168.32.130","network_if":"en0"},"logger":"Utils"} {"ts":"2024-09-22T12:50:37.085Z","level":"WARN","msg":"Set SPARK_LOCAL_IP if you need to bind to another address","logger":"Utils"} ``` example 2: SHS calls `Utils.initDaemon(log)` before loading `spark-defaults.conf`(inside construction of `HistoryServerArguments`) https://github.com/apache/spark/blob/d2e8c1cb60e34a1c7e92374c07d682aa5ca79145/core/src/main/scala/org/apache/spark/deploy/history/HistoryServer.scala#L301-L302 ``` {"ts":"2024-09-22T13:20:31.978Z","level":"INFO","msg":"Started daemon with process name: 41505H27212-MAC-01.local","logger":"HistoryServer"} {"ts":"2024-09-22T13:20:31.980Z","level":"INFO","msg":"Registering signal handler for TERM","logger":"SignalUtils"} {"ts":"2024-09-22T13:20:31.981Z","level":"INFO","msg":"Registering signal handler for HUP","logger":"SignalUtils"} {"ts":"2024-09-22T13:20:31.981Z","level":"INFO","msg":"Registering signal handler for INT","logger":"SignalUtils"} ``` then loads `spark-defaults.conf` and ignores `spark.log.structuredLogging.enabled`. ### Does this PR introduce _any_ user-facing change? No, spark structured logging is an unreleased feature. ### How was this patch tested? Write `spark.log.structuredLogging.enabled=false` in `spark-defaults.conf` 4.0.0-preview2 ``` $ SPARK_NO_DAEMONIZE=1 sbin/start-history-server.sh starting org.apache.spark.deploy.history.HistoryServer, logging to /Users/chengpan/app/spark-4.0.0-preview2-bin-hadoop3/logs/spark-chengpan-org.apache.spark.deploy.history.HistoryServer-1-H27212-MAC-01.local.out Spark Command: /Users/chengpan/.sdkman/candidates/java/current/bin/java -cp /Users/chengpan/app/spark-4.0.0-preview2-bin-hadoop3/conf/:/Users/chengpan/app/spark-4.0.0-preview2-bin-hadoop3/jars/slf4j-api-2.0.16.jar:/Users/chengpan/app/spark-4.0.0-preview2-bin-hadoop3/jars/* -Xmx1g org.apache.spark.deploy.history.HistoryServer ======================================== Using Spark's default log4j profile: org/apache/spark/log4j2-defaults.properties {"ts":"2024-09-22T12:50:37.082Z","level":"WARN","msg":"Your hostname, H27212-MAC-01.local, resolves to a loopback address: 127.0.0.1; using 192.168.32.130 instead (on interface en0)","context":{"host":"H27212-MAC-01.local","host_port":"127.0.0.1","host_port2":"192.168.32.130","network_if":"en0"},"logger":"Utils"} {"ts":"2024-09-22T12:50:37.085Z","level":"WARN","msg":"Set SPARK_LOCAL_IP if you need to bind to another address","logger":"Utils"} {"ts":"2024-09-22T12:50:37.109Z","level":"INFO","msg":"Started daemon with process name: 37764H27212-MAC-01.local","logger":"HistoryServer"} {"ts":"2024-09-22T12:50:37.112Z","level":"INFO","msg":"Registering signal handler for TERM","logger":"SignalUtils"} {"ts":"2024-09-22T12:50:37.112Z","level":"INFO","msg":"Registering signal handler for HUP","logger":"SignalUtils"} {"ts":"2024-09-22T12:50:37.112Z","level":"INFO","msg":"Registering signal handler for INT","logger":"SignalUtils"} {"ts":"2024-09-22T12:50:37.258Z","level":"WARN","msg":"Unable to load native-hadoop library for your platform... using builtin-java classes where applicable","logger":"NativeCodeLoader"} {"ts":"2024-09-22T12:50:37.275Z","level":"INFO","msg":"Changing view acls to: chengpan","logger":"SecurityManager"} {"ts":"2024-09-22T12:50:37.275Z","level":"INFO","msg":"Changing modify acls to: chengpan","logger":"SecurityManager"} {"ts":"2024-09-22T12:50:37.276Z","level":"INFO","msg":"Changing view acls groups to: chengpan","logger":"SecurityManager"} {"ts":"2024-09-22T12:50:37.276Z","level":"INFO","msg":"Changing modify acls groups to: chengpan","logger":"SecurityManager"} {"ts":"2024-09-22T12:50:37.277Z","level":"INFO","msg":"SecurityManager: authentication disabled; ui acls disabled; users with view permissions: chengpan groups with view permissions: EMPTY; users with modify permissions: chengpan; groups with modify permissions: EMPTY; RPC SSL disabled","logger":"SecurityManager"} {"ts":"2024-09-22T12:50:37.309Z","level":"INFO","msg":"History server ui acls disabled; users with admin permissions: ; groups with admin permissions: ","logger":"FsHistoryProvider"} {"ts":"2024-09-22T12:50:37.409Z","level":"INFO","msg":"Start Jetty 0.0.0.0:18080 for HistoryServerUI","logger":"JettyUtils"} {"ts":"2024-09-22T12:50:37.466Z","level":"INFO","msg":"Successfully started service 'HistoryServerUI' on port 18080.","logger":"Utils"} {"ts":"2024-09-22T12:50:37.491Z","level":"INFO","msg":"Bound HistoryServer to 0.0.0.0, and started at http://192.168.32.130:18080","logger":"HistoryServer"} ... ``` This PR ``` $ SPARK_NO_DAEMONIZE=1 sbin/start-history-server.sh starting org.apache.spark.deploy.history.HistoryServer, logging to /Users/chengpan/Projects/apache-spark/dist/logs/spark-chengpan-org.apache.spark.deploy.history.HistoryServer-1-H27212-MAC-01.local.out Spark Command: /Users/chengpan/.sdkman/candidates/java/current/bin/java -cp /Users/chengpan/Projects/apache-spark/dist/conf/:/Users/chengpan/Projects/apache-spark/dist/jars/slf4j-api-2.0.16.jar:/Users/chengpan/Projects/apache-spark/dist/jars/* -Xmx1g org.apache.spark.deploy.history.HistoryServer ======================================== Using Spark's default log4j profile: org/apache/spark/log4j2-defaults.properties {"ts":"2024-09-22T13:20:31.903Z","level":"WARN","msg":"Your hostname, H27212-MAC-01.local, resolves to a loopback address: 127.0.0.1; using 192.168.32.130 instead (on interface en0)","context":{"host":"H27212-MAC-01.local","host_port":"127.0.0.1","host_port2":"192.168.32.130","network_if":"en0"},"logger":"Utils"} {"ts":"2024-09-22T13:20:31.905Z","level":"WARN","msg":"Set SPARK_LOCAL_IP if you need to bind to another address","logger":"Utils"} {"ts":"2024-09-22T13:20:31.978Z","level":"INFO","msg":"Started daemon with process name: 41505H27212-MAC-01.local","logger":"HistoryServer"} {"ts":"2024-09-22T13:20:31.980Z","level":"INFO","msg":"Registering signal handler for TERM","logger":"SignalUtils"} {"ts":"2024-09-22T13:20:31.981Z","level":"INFO","msg":"Registering signal handler for HUP","logger":"SignalUtils"} {"ts":"2024-09-22T13:20:31.981Z","level":"INFO","msg":"Registering signal handler for INT","logger":"SignalUtils"} {"ts":"2024-09-22T13:20:32.136Z","level":"WARN","msg":"Unable to load native-hadoop library for your platform... using builtin-java classes where applicable","logger":"NativeCodeLoader"} Using Spark's default log4j profile: org/apache/spark/log4j2-pattern-layout-defaults.properties 24/09/22 21:20:32 INFO SecurityManager: Changing view acls to: chengpan 24/09/22 21:20:32 INFO SecurityManager: Changing modify acls to: chengpan 24/09/22 21:20:32 INFO SecurityManager: Changing view acls groups to: chengpan 24/09/22 21:20:32 INFO SecurityManager: Changing modify acls groups to: chengpan 24/09/22 21:20:32 INFO SecurityManager: SecurityManager: authentication disabled; ui acls disabled; users with view permissions: chengpan groups with view permissions: EMPTY; users with modify permissions: chengpan; groups with modify permissions: EMPTY; RPC SSL disabled 24/09/22 21:20:32 INFO FsHistoryProvider: History server ui acls disabled; users with admin permissions: ; groups with admin permissions: 24/09/22 21:20:32 INFO JettyUtils: Start Jetty 0.0.0.0:18080 for HistoryServerUI 24/09/22 21:20:32 INFO Utils: Successfully started service 'HistoryServerUI' on port 18080. 24/09/22 21:20:32 INFO HistoryServer: Bound HistoryServer to 0.0.0.0, and started at http://192.168.32.130:18080 ... ``` ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48198 from pan3793/SPARK-49732. Authored-by: Cheng Pan Signed-off-by: Dongjoon Hyun --- .../spark/deploy/ExternalShuffleService.scala | 3 +++ .../scala/org/apache/spark/deploy/SparkSubmit.scala | 6 +----- .../deploy/history/HistoryServerArguments.scala | 3 +++ .../spark/deploy/master/MasterArguments.scala | 3 +++ .../spark/deploy/worker/WorkerArguments.scala | 4 ++++ .../executor/CoarseGrainedExecutorBackend.scala | 4 ++++ .../main/scala/org/apache/spark/util/Utils.scala | 13 +++++++++++++ .../cluster/k8s/KubernetesExecutorBackend.scala | 4 ++++ 8 files changed, 35 insertions(+), 5 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala index f0dcf344ce0da..57b0647e59fd9 100644 --- a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala +++ b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala @@ -169,6 +169,9 @@ object ExternalShuffleService extends Logging { Utils.initDaemon(log) val sparkConf = new SparkConf Utils.loadDefaultSparkProperties(sparkConf) + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(sparkConf) + Logging.uninitialize() val securityManager = new SecurityManager(sparkConf) // we override this value since this service is started from the command line diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala index f3833e85a482e..85ed441d58fd1 100644 --- a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala +++ b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala @@ -79,11 +79,7 @@ private[spark] class SparkSubmit extends Logging { } else { // For non-shell applications, enable structured logging if it's not explicitly disabled // via the configuration `spark.log.structuredLogging.enabled`. - if (sparkConf.getBoolean(STRUCTURED_LOGGING_ENABLED.key, defaultValue = true)) { - Logging.enableStructuredLogging() - } else { - Logging.disableStructuredLogging() - } + Utils.resetStructuredLogging(sparkConf) } // We should initialize log again after `spark.log.structuredLogging.enabled` effected diff --git a/core/src/main/scala/org/apache/spark/deploy/history/HistoryServerArguments.scala b/core/src/main/scala/org/apache/spark/deploy/history/HistoryServerArguments.scala index 2fdf7a473a298..f1343a0551384 100644 --- a/core/src/main/scala/org/apache/spark/deploy/history/HistoryServerArguments.scala +++ b/core/src/main/scala/org/apache/spark/deploy/history/HistoryServerArguments.scala @@ -53,6 +53,9 @@ private[history] class HistoryServerArguments(conf: SparkConf, args: Array[Strin // This mutates the SparkConf, so all accesses to it must be made after this line Utils.loadDefaultSparkProperties(conf, propertiesFile) + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(conf) + Logging.uninitialize() // scalastyle:off line.size.limit println private def printUsageAndExit(exitCode: Int, error: String = ""): Unit = { diff --git a/core/src/main/scala/org/apache/spark/deploy/master/MasterArguments.scala b/core/src/main/scala/org/apache/spark/deploy/master/MasterArguments.scala index 045a3da74dcd0..6647b11874d72 100644 --- a/core/src/main/scala/org/apache/spark/deploy/master/MasterArguments.scala +++ b/core/src/main/scala/org/apache/spark/deploy/master/MasterArguments.scala @@ -53,6 +53,9 @@ private[master] class MasterArguments(args: Array[String], conf: SparkConf) exte // This mutates the SparkConf, so all accesses to it must be made after this line propertiesFile = Utils.loadDefaultSparkProperties(conf, propertiesFile) + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(conf) + Logging.uninitialize() if (conf.contains(MASTER_UI_PORT.key)) { webUiPort = conf.get(MASTER_UI_PORT) diff --git a/core/src/main/scala/org/apache/spark/deploy/worker/WorkerArguments.scala b/core/src/main/scala/org/apache/spark/deploy/worker/WorkerArguments.scala index 94a27e1a3e6da..f24cd59418300 100644 --- a/core/src/main/scala/org/apache/spark/deploy/worker/WorkerArguments.scala +++ b/core/src/main/scala/org/apache/spark/deploy/worker/WorkerArguments.scala @@ -22,6 +22,7 @@ import java.lang.management.ManagementFactory import scala.annotation.tailrec import org.apache.spark.SparkConf +import org.apache.spark.internal.Logging import org.apache.spark.internal.config.Worker._ import org.apache.spark.util.{IntParam, MemoryParam, Utils} @@ -59,6 +60,9 @@ private[worker] class WorkerArguments(args: Array[String], conf: SparkConf) { // This mutates the SparkConf, so all accesses to it must be made after this line propertiesFile = Utils.loadDefaultSparkProperties(conf, propertiesFile) + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(conf) + Logging.uninitialize() conf.get(WORKER_UI_PORT).foreach { webUiPort = _ } diff --git a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala index eaa07b9a81f5b..e880cf8da9ec2 100644 --- a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala +++ b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala @@ -468,6 +468,10 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging { } } + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(driverConf) + Logging.uninitialize() + cfg.hadoopDelegationCreds.foreach { tokens => SparkHadoopUtil.get.addDelegationTokens(tokens, driverConf) } diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala index 52213f36a2cd1..5703128aacbb9 100644 --- a/core/src/main/scala/org/apache/spark/util/Utils.scala +++ b/core/src/main/scala/org/apache/spark/util/Utils.scala @@ -2673,6 +2673,19 @@ private[spark] object Utils } } + /** + * Utility function to enable or disable structured logging based on SparkConf. + * This is designed for a code path which logging system may be initilized before + * loading SparkConf. + */ + def resetStructuredLogging(sparkConf: SparkConf): Unit = { + if (sparkConf.getBoolean(STRUCTURED_LOGGING_ENABLED.key, defaultValue = true)) { + Logging.enableStructuredLogging() + } else { + Logging.disableStructuredLogging() + } + } + /** * Return the jar files pointed by the "spark.jars" property. Spark internally will distribute * these jars through file server. In the YARN mode, it will return an empty list, since YARN diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala index c515ae5e3a246..e44d7e29ef606 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala @@ -116,6 +116,10 @@ private[spark] object KubernetesExecutorBackend extends Logging { } } + // Initialize logging system again after `spark.log.structuredLogging.enabled` takes effect + Utils.resetStructuredLogging(driverConf) + Logging.uninitialize() + cfg.hadoopDelegationCreds.foreach { tokens => SparkHadoopUtil.get.addDelegationTokens(tokens, driverConf) } From da106f86260b8138df7c5da5e05af9c801fc318d Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Mon, 30 Sep 2024 23:33:06 -0700 Subject: [PATCH 181/189] [SPARK-49840][INFRA] Use `MacOS 15` in `build_maven_java21_macos14.yml` ### What changes were proposed in this pull request? This PR aims to upgrade `MacOS` from `14` to `15` in `build_maven_java21_macos14.yml`. ### Why are the changes needed? To use the latest MacOS as a part of Apache Spark 4.0.0 preparation. - https://github.com/actions/runner-images/blob/main/images/macos/macos-15-arm64-Readme.md ### Does this PR introduce _any_ user-facing change? No. This is an infra change. ### How was this patch tested? N/A. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48305 from dongjoon-hyun/SPARK-49840. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- ...aven_java21_macos14.yml => build_maven_java21_macos15.yml} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename .github/workflows/{build_maven_java21_macos14.yml => build_maven_java21_macos15.yml} (92%) diff --git a/.github/workflows/build_maven_java21_macos14.yml b/.github/workflows/build_maven_java21_macos15.yml similarity index 92% rename from .github/workflows/build_maven_java21_macos14.yml rename to .github/workflows/build_maven_java21_macos15.yml index fb5e609f4eae0..cc6d0ea4e90da 100644 --- a/.github/workflows/build_maven_java21_macos14.yml +++ b/.github/workflows/build_maven_java21_macos15.yml @@ -17,7 +17,7 @@ # under the License. # -name: "Build / Maven (master, Scala 2.13, Hadoop 3, JDK 21, macos-14)" +name: "Build / Maven (master, Scala 2.13, Hadoop 3, JDK 21, MacOS-15)" on: schedule: @@ -32,7 +32,7 @@ jobs: if: github.repository == 'apache/spark' with: java: 21 - os: macos-14 + os: macos-15 envs: >- { "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES" From 8d0f6fb902219adfa5dd019a88c5ef4e8bf2ed7c Mon Sep 17 00:00:00 2001 From: panbingkun Date: Mon, 30 Sep 2024 23:36:39 -0700 Subject: [PATCH 182/189] [SPARK-49826][BUILD] Upgrade jackson to 2.18.0 ### What changes were proposed in this pull request? The pr aims to upgrade `jackson` from `2.17.2` to `2.18.0` ### Why are the changes needed? The full release notes: https://github.com/FasterXML/jackson/wiki/Jackson-Release-2.18.0 image ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48294 from panbingkun/SPARK-49826. Authored-by: panbingkun Signed-off-by: Dongjoon Hyun --- dev/deps/spark-deps-hadoop-3-hive-2.3 | 14 +++++++------- pom.xml | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 95a667ccfc72d..f6ce3d25ebc8a 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -105,16 +105,16 @@ ini4j/0.5.4//ini4j-0.5.4.jar istack-commons-runtime/3.0.8//istack-commons-runtime-3.0.8.jar ivy/2.5.2//ivy-2.5.2.jar j2objc-annotations/3.0.0//j2objc-annotations-3.0.0.jar -jackson-annotations/2.17.2//jackson-annotations-2.17.2.jar +jackson-annotations/2.18.0//jackson-annotations-2.18.0.jar jackson-core-asl/1.9.13//jackson-core-asl-1.9.13.jar -jackson-core/2.17.2//jackson-core-2.17.2.jar -jackson-databind/2.17.2//jackson-databind-2.17.2.jar -jackson-dataformat-cbor/2.17.2//jackson-dataformat-cbor-2.17.2.jar -jackson-dataformat-yaml/2.17.2//jackson-dataformat-yaml-2.17.2.jar +jackson-core/2.18.0//jackson-core-2.18.0.jar +jackson-databind/2.18.0//jackson-databind-2.18.0.jar +jackson-dataformat-cbor/2.18.0//jackson-dataformat-cbor-2.18.0.jar +jackson-dataformat-yaml/2.18.0//jackson-dataformat-yaml-2.18.0.jar jackson-datatype-jdk8/2.17.0//jackson-datatype-jdk8-2.17.0.jar -jackson-datatype-jsr310/2.17.2//jackson-datatype-jsr310-2.17.2.jar +jackson-datatype-jsr310/2.18.0//jackson-datatype-jsr310-2.18.0.jar jackson-mapper-asl/1.9.13//jackson-mapper-asl-1.9.13.jar -jackson-module-scala_2.13/2.17.2//jackson-module-scala_2.13-2.17.2.jar +jackson-module-scala_2.13/2.18.0//jackson-module-scala_2.13-2.18.0.jar jakarta.annotation-api/2.0.0//jakarta.annotation-api-2.0.0.jar jakarta.inject-api/2.0.1//jakarta.inject-api-2.0.1.jar jakarta.servlet-api/5.0.0//jakarta.servlet-api-5.0.0.jar diff --git a/pom.xml b/pom.xml index 4bdb92d86a727..6a77da703dbd2 100644 --- a/pom.xml +++ b/pom.xml @@ -180,8 +180,8 @@ true true 1.9.13 - 2.17.2 - 2.17.2 + 2.18.0 + 2.18.0 2.3.1 1.1.10.7 3.0.3 From c0a1ea2a4c4218fc15b8f990ed2f5ea99755d322 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 30 Sep 2024 23:45:21 -0700 Subject: [PATCH 183/189] [SPARK-49795][CORE][SQL][SS][DSTREAM][ML][MLLIB][K8S][YARN][EXAMPLES] Clean up deprecated Guava API usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? In order to clean up the usage of deprecated Guava API, the following changes were made in this pr: 1. Replaced `Files.write(from, to, charset)` with `Files.asCharSink(to, charset).write(from)`. This change was made with reference to: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/io/Files.java#L275-L291 ```java /** * Writes a character sequence (such as a string) to a file using the given character set. * * param from the character sequence to write * param to the destination file * param charset the charset used to encode the output stream; see {link StandardCharsets} for * helpful predefined constants * throws IOException if an I/O error occurs * deprecated Prefer {code asCharSink(to, charset).write(from)}. */ Deprecated InlineMe( replacement = "Files.asCharSink(to, charset).write(from)", imports = "com.google.common.io.Files") public static void write(CharSequence from, File to, Charset charset) throws IOException { asCharSink(to, charset).write(from); } ``` 2. Replaced `Files.append(from, to, charset)` with `Files.asCharSink(to, charset, FileWriteMode.APPEND).write(from)`. This change was made with reference to: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/io/Files.java#L350-L368 ```java /** * Appends a character sequence (such as a string) to a file using the given character set. * * param from the character sequence to append * param to the destination file * param charset the charset used to encode the output stream; see {link StandardCharsets} for * helpful predefined constants * throws IOException if an I/O error occurs * deprecated Prefer {code asCharSink(to, charset, FileWriteMode.APPEND).write(from)}. This * method is scheduled to be removed in October 2019. */ Deprecated InlineMe( replacement = "Files.asCharSink(to, charset, FileWriteMode.APPEND).write(from)", imports = {"com.google.common.io.FileWriteMode", "com.google.common.io.Files"}) public static void append(CharSequence from, File to, Charset charset) throws IOException { asCharSink(to, charset, FileWriteMode.APPEND).write(from); } ``` 3. Replaced `Files.toString(file, charset)` with `Files.asCharSource(file, charset).read()`. This change was made with reference to: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/io/Files.java#L243-L259 ```java /** * Reads all characters from a file into a {link String}, using the given character set. * * param file the file to read from * param charset the charset used to decode the input stream; see {link StandardCharsets} for * helpful predefined constants * return a string containing all the characters from the file * throws IOException if an I/O error occurs * deprecated Prefer {code asCharSource(file, charset).read()}. */ Deprecated InlineMe( replacement = "Files.asCharSource(file, charset).read()", imports = "com.google.common.io.Files") public static String toString(File file, Charset charset) throws IOException { return asCharSource(file, charset).read(); } ``` 4. Replaced `HashFunction.murmur3_32()` with `HashFunction.murmur3_32_fixed()`. This change was made with reference to: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/hash/Hashing.java#L99-L115 ```java /** * Returns a hash function implementing the 32-bit murmur3 * algorithm, x86 variant (little-endian variant), using the given seed value, with a known * bug as described in the deprecation text. * *

    The C++ equivalent is the MurmurHash3_x86_32 function (Murmur3A), which however does not * have the bug. * * deprecated This implementation produces incorrect hash values from the {link * HashFunction#hashString} method if the string contains non-BMP characters. Use {link * #murmur3_32_fixed()} instead. */ Deprecated public static HashFunction murmur3_32() { return Murmur3_32HashFunction.MURMUR3_32; } ``` This change is safe for Spark. The difference between `MURMUR3_32` and `MURMUR3_32_FIXED` lies in the different `supplementaryPlaneFix` parameters passed when constructing the `Murmur3_32HashFunction`: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/hash/Murmur3_32HashFunction.java#L56-L59 ```java static final HashFunction MURMUR3_32 = new Murmur3_32HashFunction(0, /* supplementaryPlaneFix= */ false); static final HashFunction MURMUR3_32_FIXED = new Murmur3_32HashFunction(0, /* supplementaryPlaneFix= */ true); ``` However, the `supplementaryPlaneFix` parameter is only used in `Murmur3_32HashFunction#hashString`, and Spark only utilizes `Murmur3_32HashFunction#hashInt`. Therefore, there will be no logical changes to this method after this change. https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/hash/Murmur3_32HashFunction.java#L108-L114 ```java Override public HashCode hashInt(int input) { int k1 = mixK1(input); int h1 = mixH1(seed, k1); return fmix(h1, Ints.BYTES); } ``` 5. Replaced `Throwables.propagateIfPossible(throwable, declaredType)` with `Throwables.throwIfInstanceOf(throwable, declaredType)` + `Throwables.throwIfUnchecked(throwable)`. This change was made with reference to: https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/base/Throwables.java#L156-L175 ``` /** * Propagates {code throwable} exactly as-is, if and only if it is an instance of {link * RuntimeException}, {link Error}, or {code declaredType}. * *

    Discouraged in favor of calling {link #throwIfInstanceOf} and {link * #throwIfUnchecked}. * * param throwable the Throwable to possibly propagate * param declaredType the single checked exception type declared by the calling method * deprecated Use a combination of {link #throwIfInstanceOf} and {link #throwIfUnchecked}, * which togther provide the same behavior except that they reject {code null}. */ Deprecated J2ktIncompatible GwtIncompatible // propagateIfInstanceOf public static void propagateIfPossible( CheckForNull Throwable throwable, Class declaredType) throws X { propagateIfInstanceOf(throwable, declaredType); propagateIfPossible(throwable); } ``` 6. Made modifications to `Throwables.propagate` with reference to https://github.com/google/guava/wiki/Why-we-deprecated-Throwables.propagate - For cases where it is known to be a checked exception, including `IOException`, `GeneralSecurityException`, `SaslException`, and `RocksDBException`, none of which are subclasses of `RuntimeException` or `Error`, directly replaced `Throwables.propagate(e)` with `throw new RuntimeException(e);`. - For cases where it cannot be determined whether it is a checked exception or an unchecked exception or Error, use ```java throwIfUnchecked(e); throw new RuntimeException(e); ``` to replace `Throwables.propagate(e)`。 https://github.com/google/guava/blob/0c33dd12b193402cdf6962d43d69743521aa2f76/guava/src/com/google/common/base/Throwables.java#L199-L235 ```java /** * ... * deprecated To preserve behavior, use {code throw e} or {code throw new RuntimeException(e)} * directly, or use a combination of {link #throwIfUnchecked} and {code throw new * RuntimeException(e)}. But consider whether users would be better off if your API threw a * different type of exception. For background on the deprecation, read Why we deprecated {code Throwables.propagate}. */ CanIgnoreReturnValue J2ktIncompatible GwtIncompatible Deprecated public static RuntimeException propagate(Throwable throwable) { throwIfUnchecked(throwable); throw new RuntimeException(throwable); } ``` ### Why are the changes needed? Clean up deprecated Guava API usage. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Pass GitHub Actions ### Was this patch authored or co-authored using generative AI tooling? No Closes #48248 from LuciferYang/guava-deprecation. Authored-by: yangjie01 Signed-off-by: Dongjoon Hyun --- .../apache/spark/util/kvstore/LevelDB.java | 3 +- .../spark/util/kvstore/LevelDBIterator.java | 5 +-- .../apache/spark/util/kvstore/RocksDB.java | 3 +- .../spark/util/kvstore/RocksDBIterator.java | 5 +-- .../spark/network/client/TransportClient.java | 6 ++-- .../client/TransportClientFactory.java | 3 +- .../network/crypto/AuthClientBootstrap.java | 3 +- .../spark/network/crypto/AuthRpcHandler.java | 3 +- .../spark/network/sasl/SparkSaslClient.java | 7 ++--- .../spark/network/sasl/SparkSaslServer.java | 5 ++- .../network/shuffledb/LevelDBIterator.java | 4 +-- .../spark/network/shuffledb/RocksDB.java | 7 ++--- .../network/shuffledb/RocksDBIterator.java | 3 +- .../spark/sql/kafka010/KafkaTestUtils.scala | 4 +-- .../apache/spark/io/ReadAheadInputStream.java | 3 +- .../scala/org/apache/spark/TestUtils.scala | 2 +- .../spark/deploy/worker/DriverRunner.scala | 4 +-- .../spark/deploy/worker/ExecutorRunner.scala | 2 +- .../spark/util/collection/AppendOnlyMap.scala | 2 +- .../spark/util/collection/OpenHashSet.scala | 2 +- .../test/org/apache/spark/JavaAPISuite.java | 2 +- .../scala/org/apache/spark/FileSuite.scala | 4 +-- .../org/apache/spark/SparkContextSuite.scala | 31 ++++++++++--------- .../history/EventLogFileReadersSuite.scala | 6 ++-- .../history/FsHistoryProviderSuite.scala | 3 +- .../history/HistoryServerArgumentsSuite.scala | 4 +-- .../deploy/history/HistoryServerSuite.scala | 2 +- .../plugin/PluginContainerSuite.scala | 2 +- .../ResourceDiscoveryPluginSuite.scala | 2 +- .../org/apache/spark/rpc/RpcEnvSuite.scala | 12 +++---- .../apache/spark/util/FileAppenderSuite.scala | 6 ++-- .../org/apache/spark/util/UtilsSuite.scala | 6 ++-- .../JavaRecoverableNetworkWordCount.java | 4 ++- .../RecoverableNetworkWordCount.scala | 5 +-- .../libsvm/JavaLibSVMRelationSuite.java | 2 +- .../source/libsvm/LibSVMRelationSuite.scala | 6 ++-- .../spark/mllib/util/MLUtilsSuite.scala | 6 ++-- .../k8s/SparkKubernetesClientFactory.scala | 2 +- .../HadoopConfDriverFeatureStep.scala | 2 +- .../KerberosConfDriverFeatureStep.scala | 2 +- .../features/PodTemplateConfigMapStep.scala | 2 +- ...ubernetesCredentialsFeatureStepSuite.scala | 2 +- .../HadoopConfDriverFeatureStepSuite.scala | 2 +- .../HadoopConfExecutorFeatureStepSuite.scala | 2 +- .../KerberosConfDriverFeatureStepSuite.scala | 4 +-- .../integrationtest/DecommissionSuite.scala | 6 ++-- .../k8s/integrationtest/KubernetesSuite.scala | 2 +- .../deploy/yarn/BaseYarnClusterSuite.scala | 6 ++-- .../spark/deploy/yarn/YarnClusterSuite.scala | 29 +++++++++-------- .../yarn/YarnShuffleIntegrationSuite.scala | 2 +- .../arrow/ArrowConvertersSuite.scala | 6 ++-- .../HiveThriftServer2Suites.scala | 6 ++-- .../hive/thriftserver/UISeleniumSuite.scala | 6 ++-- .../sql/hive/execution/SQLQuerySuite.scala | 17 +++++----- .../apache/spark/streaming/JavaAPISuite.java | 2 +- .../spark/streaming/CheckpointSuite.scala | 2 +- .../spark/streaming/InputStreamsSuite.scala | 10 +++--- .../spark/streaming/MasterFailureTest.scala | 2 +- 58 files changed, 148 insertions(+), 145 deletions(-) diff --git a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDB.java b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDB.java index 13a9d89f4705c..7f8d6c58aec7e 100644 --- a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDB.java +++ b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDB.java @@ -255,7 +255,8 @@ public Iterator iterator() { iteratorTracker.add(new WeakReference<>(it)); return it; } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } }; diff --git a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDBIterator.java b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDBIterator.java index 69757fdc65d68..29ed37ffa44e5 100644 --- a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDBIterator.java +++ b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDBIterator.java @@ -127,7 +127,7 @@ public boolean hasNext() { try { close(); } catch (IOException ioe) { - throw Throwables.propagate(ioe); + throw new RuntimeException(ioe); } } return next != null; @@ -151,7 +151,8 @@ public T next() { next = null; return ret; } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } diff --git a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDB.java b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDB.java index dc7ad0be5c007..4bc2b233fe12d 100644 --- a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDB.java +++ b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDB.java @@ -287,7 +287,8 @@ public Iterator iterator() { iteratorTracker.add(new WeakReference<>(it)); return it; } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } }; diff --git a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDBIterator.java b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDBIterator.java index a98b0482e35cc..e350ddc2d445a 100644 --- a/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDBIterator.java +++ b/common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDBIterator.java @@ -113,7 +113,7 @@ public boolean hasNext() { try { close(); } catch (IOException ioe) { - throw Throwables.propagate(ioe); + throw new RuntimeException(ioe); } } return next != null; @@ -137,7 +137,8 @@ public T next() { next = null; return ret; } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java index 4c144a73a9299..a9df47645d36f 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java +++ b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java @@ -290,9 +290,11 @@ public void onFailure(Throwable e) { try { return result.get(timeoutMs, TimeUnit.MILLISECONDS); } catch (ExecutionException e) { - throw Throwables.propagate(e.getCause()); + Throwables.throwIfUnchecked(e.getCause()); + throw new RuntimeException(e.getCause()); } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java index e1f19f956cc0a..d64b8c8f838e9 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java +++ b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java @@ -342,7 +342,8 @@ public void operationComplete(final Future handshakeFuture) { logger.error("Exception while bootstrapping client after {} ms", e, MDC.of(LogKeys.BOOTSTRAP_TIME$.MODULE$, bootstrapTimeMs)); client.close(); - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } long postBootstrap = System.nanoTime(); diff --git a/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthClientBootstrap.java b/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthClientBootstrap.java index 08e2c084fe67b..2e9ccd0e0ad21 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthClientBootstrap.java +++ b/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthClientBootstrap.java @@ -22,7 +22,6 @@ import java.security.GeneralSecurityException; import java.util.concurrent.TimeoutException; -import com.google.common.base.Throwables; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.Channel; @@ -80,7 +79,7 @@ public void doBootstrap(TransportClient client, Channel channel) { doSparkAuth(client, channel); client.setClientId(appId); } catch (GeneralSecurityException | IOException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } catch (RuntimeException e) { // There isn't a good exception that can be caught here to know whether it's really // OK to switch back to SASL (because the server doesn't speak the new protocol). So diff --git a/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthRpcHandler.java b/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthRpcHandler.java index 65367743e24f9..087e3d21e22bb 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthRpcHandler.java +++ b/common/network-common/src/main/java/org/apache/spark/network/crypto/AuthRpcHandler.java @@ -132,7 +132,8 @@ protected boolean doAuthChallenge( try { engine.close(); } catch (Exception e) { - throw Throwables.propagate(e); + Throwables.throwIfUnchecked(e); + throw new RuntimeException(e); } } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslClient.java b/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslClient.java index 3600c1045dbf4..a61b1c3c0c416 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslClient.java +++ b/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslClient.java @@ -29,7 +29,6 @@ import javax.security.sasl.SaslClient; import javax.security.sasl.SaslException; -import com.google.common.base.Throwables; import com.google.common.collect.ImmutableMap; import org.apache.spark.internal.SparkLogger; @@ -62,7 +61,7 @@ public SparkSaslClient(String secretKeyId, SecretKeyHolder secretKeyHolder, bool this.saslClient = Sasl.createSaslClient(new String[] { DIGEST }, null, null, DEFAULT_REALM, saslProps, new ClientCallbackHandler()); } catch (SaslException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } @@ -72,7 +71,7 @@ public synchronized byte[] firstToken() { try { return saslClient.evaluateChallenge(new byte[0]); } catch (SaslException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } else { return new byte[0]; @@ -98,7 +97,7 @@ public synchronized byte[] response(byte[] token) { try { return saslClient != null ? saslClient.evaluateChallenge(token) : new byte[0]; } catch (SaslException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslServer.java b/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslServer.java index b897650afe832..f32fd5145c7c5 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslServer.java +++ b/common/network-common/src/main/java/org/apache/spark/network/sasl/SparkSaslServer.java @@ -31,7 +31,6 @@ import java.util.Map; import com.google.common.base.Preconditions; -import com.google.common.base.Throwables; import com.google.common.collect.ImmutableMap; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; @@ -94,7 +93,7 @@ public SparkSaslServer( this.saslServer = Sasl.createSaslServer(DIGEST, null, DEFAULT_REALM, saslProps, new DigestCallbackHandler()); } catch (SaslException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } @@ -119,7 +118,7 @@ public synchronized byte[] response(byte[] token) { try { return saslServer != null ? saslServer.evaluateResponse(token) : new byte[0]; } catch (SaslException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/LevelDBIterator.java b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/LevelDBIterator.java index 5796e34a6f05e..2ac549775449a 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/LevelDBIterator.java +++ b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/LevelDBIterator.java @@ -17,8 +17,6 @@ package org.apache.spark.network.shuffledb; -import com.google.common.base.Throwables; - import java.io.IOException; import java.util.Map; import java.util.NoSuchElementException; @@ -47,7 +45,7 @@ public boolean hasNext() { try { close(); } catch (IOException ioe) { - throw Throwables.propagate(ioe); + throw new RuntimeException(ioe); } } return next != null; diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDB.java b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDB.java index d33895d6c2d62..2737ab8ed754c 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDB.java +++ b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDB.java @@ -19,7 +19,6 @@ import java.io.IOException; -import com.google.common.base.Throwables; import org.rocksdb.RocksDBException; /** @@ -37,7 +36,7 @@ public void put(byte[] key, byte[] value) { try { db.put(key, value); } catch (RocksDBException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } @@ -46,7 +45,7 @@ public byte[] get(byte[] key) { try { return db.get(key); } catch (RocksDBException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } @@ -55,7 +54,7 @@ public void delete(byte[] key) { try { db.delete(key); } catch (RocksDBException e) { - throw Throwables.propagate(e); + throw new RuntimeException(e); } } diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDBIterator.java b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDBIterator.java index 78562f91a4b75..829a7ded6330b 100644 --- a/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDBIterator.java +++ b/common/network-common/src/main/java/org/apache/spark/network/shuffledb/RocksDBIterator.java @@ -22,7 +22,6 @@ import java.util.Map; import java.util.NoSuchElementException; -import com.google.common.base.Throwables; import org.rocksdb.RocksIterator; /** @@ -52,7 +51,7 @@ public boolean hasNext() { try { close(); } catch (IOException ioe) { - throw Throwables.propagate(ioe); + throw new RuntimeException(ioe); } } return next != null; diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala index 7852bc814ccd4..c3f02eebab23a 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala @@ -176,7 +176,7 @@ class KafkaTestUtils( } kdc.getKrb5conf.delete() - Files.write(krb5confStr, kdc.getKrb5conf, StandardCharsets.UTF_8) + Files.asCharSink(kdc.getKrb5conf, StandardCharsets.UTF_8).write(krb5confStr) logDebug(s"krb5.conf file content: $krb5confStr") } @@ -240,7 +240,7 @@ class KafkaTestUtils( | principal="$kafkaServerUser@$realm"; |}; """.stripMargin.trim - Files.write(content, file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(content) logDebug(s"Created JAAS file: ${file.getPath}") logDebug(s"JAAS file content: $content") file.getAbsolutePath() diff --git a/core/src/main/java/org/apache/spark/io/ReadAheadInputStream.java b/core/src/main/java/org/apache/spark/io/ReadAheadInputStream.java index 5e9f1b78273a5..7dd87df713e6e 100644 --- a/core/src/main/java/org/apache/spark/io/ReadAheadInputStream.java +++ b/core/src/main/java/org/apache/spark/io/ReadAheadInputStream.java @@ -120,7 +120,8 @@ private boolean isEndOfStream() { private void checkReadException() throws IOException { if (readAborted) { - Throwables.propagateIfPossible(readException, IOException.class); + Throwables.throwIfInstanceOf(readException, IOException.class); + Throwables.throwIfUnchecked(readException); throw new IOException(readException); } } diff --git a/core/src/main/scala/org/apache/spark/TestUtils.scala b/core/src/main/scala/org/apache/spark/TestUtils.scala index 5e3078d7292ba..fed15a067c00f 100644 --- a/core/src/main/scala/org/apache/spark/TestUtils.scala +++ b/core/src/main/scala/org/apache/spark/TestUtils.scala @@ -421,7 +421,7 @@ private[spark] object TestUtils extends SparkTestUtils { def createTempScriptWithExpectedOutput(dir: File, prefix: String, output: String): String = { val file = File.createTempFile(prefix, ".sh", dir) val script = s"cat < expected = Arrays.asList("1", "2", "3", "4"); diff --git a/core/src/test/scala/org/apache/spark/FileSuite.scala b/core/src/test/scala/org/apache/spark/FileSuite.scala index 5651dc9b2dbdc..5f9912cbd021d 100644 --- a/core/src/test/scala/org/apache/spark/FileSuite.scala +++ b/core/src/test/scala/org/apache/spark/FileSuite.scala @@ -334,8 +334,8 @@ class FileSuite extends SparkFunSuite with LocalSparkContext { for (i <- 0 until 8) { val tempFile = new File(tempDir, s"part-0000$i") - Files.write("someline1 in file1\nsomeline2 in file1\nsomeline3 in file1", tempFile, - StandardCharsets.UTF_8) + Files.asCharSink(tempFile, StandardCharsets.UTF_8) + .write("someline1 in file1\nsomeline2 in file1\nsomeline3 in file1") } for (p <- Seq(1, 2, 8)) { diff --git a/core/src/test/scala/org/apache/spark/SparkContextSuite.scala b/core/src/test/scala/org/apache/spark/SparkContextSuite.scala index 12f9d2f83c777..44b2da603a1f6 100644 --- a/core/src/test/scala/org/apache/spark/SparkContextSuite.scala +++ b/core/src/test/scala/org/apache/spark/SparkContextSuite.scala @@ -119,8 +119,8 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu val absolutePath2 = file2.getAbsolutePath try { - Files.write("somewords1", file1, StandardCharsets.UTF_8) - Files.write("somewords2", file2, StandardCharsets.UTF_8) + Files.asCharSink(file1, StandardCharsets.UTF_8).write("somewords1") + Files.asCharSink(file2, StandardCharsets.UTF_8).write("somewords2") val length1 = file1.length() val length2 = file2.length() @@ -178,10 +178,10 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu s"${jarFile.getParent}/../${jarFile.getParentFile.getName}/${jarFile.getName}#zoo" try { - Files.write("somewords1", file1, StandardCharsets.UTF_8) - Files.write("somewords22", file2, StandardCharsets.UTF_8) - Files.write("somewords333", file3, StandardCharsets.UTF_8) - Files.write("somewords4444", file4, StandardCharsets.UTF_8) + Files.asCharSink(file1, StandardCharsets.UTF_8).write("somewords1") + Files.asCharSink(file2, StandardCharsets.UTF_8).write("somewords22") + Files.asCharSink(file3, StandardCharsets.UTF_8).write("somewords333") + Files.asCharSink(file4, StandardCharsets.UTF_8).write("somewords4444") val length1 = file1.length() val length2 = file2.length() val length3 = file1.length() @@ -373,8 +373,8 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu assert(subdir2.mkdir()) val file1 = new File(subdir1, "file") val file2 = new File(subdir2, "file") - Files.write("old", file1, StandardCharsets.UTF_8) - Files.write("new", file2, StandardCharsets.UTF_8) + Files.asCharSink(file1, StandardCharsets.UTF_8).write("old") + Files.asCharSink(file2, StandardCharsets.UTF_8).write("new") sc = new SparkContext("local-cluster[1,1,1024]", "test") sc.addFile(file1.getAbsolutePath) def getAddedFileContents(): String = { @@ -503,12 +503,15 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu try { // Create 5 text files. - Files.write("someline1 in file1\nsomeline2 in file1\nsomeline3 in file1", file1, - StandardCharsets.UTF_8) - Files.write("someline1 in file2\nsomeline2 in file2", file2, StandardCharsets.UTF_8) - Files.write("someline1 in file3", file3, StandardCharsets.UTF_8) - Files.write("someline1 in file4\nsomeline2 in file4", file4, StandardCharsets.UTF_8) - Files.write("someline1 in file2\nsomeline2 in file5", file5, StandardCharsets.UTF_8) + Files.asCharSink(file1, StandardCharsets.UTF_8) + .write("someline1 in file1\nsomeline2 in file1\nsomeline3 in file1") + Files.asCharSink(file2, StandardCharsets.UTF_8) + .write("someline1 in file2\nsomeline2 in file2") + Files.asCharSink(file3, StandardCharsets.UTF_8).write("someline1 in file3") + Files.asCharSink(file4, StandardCharsets.UTF_8) + .write("someline1 in file4\nsomeline2 in file4") + Files.asCharSink(file5, StandardCharsets.UTF_8) + .write("someline1 in file2\nsomeline2 in file5") sc = new SparkContext(new SparkConf().setAppName("test").setMaster("local")) diff --git a/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileReadersSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileReadersSuite.scala index f34f792881f90..7501a98a1a573 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileReadersSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileReadersSuite.scala @@ -221,7 +221,7 @@ class SingleFileEventLogFileReaderSuite extends EventLogFileReadersSuite { val entry = is.getNextEntry assert(entry != null) val actual = new String(ByteStreams.toByteArray(is), StandardCharsets.UTF_8) - val expected = Files.toString(new File(logPath.toString), StandardCharsets.UTF_8) + val expected = Files.asCharSource(new File(logPath.toString), StandardCharsets.UTF_8).read() assert(actual === expected) assert(is.getNextEntry === null) } @@ -368,8 +368,8 @@ class RollingEventLogFilesReaderSuite extends EventLogFileReadersSuite { assert(allFileNames.contains(fileName)) val actual = new String(ByteStreams.toByteArray(is), StandardCharsets.UTF_8) - val expected = Files.toString(new File(logPath.toString, fileName), - StandardCharsets.UTF_8) + val expected = Files.asCharSource( + new File(logPath.toString, fileName), StandardCharsets.UTF_8).read() assert(actual === expected) } } diff --git a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala index 3013a5bf4a294..852f94bda870d 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala @@ -708,7 +708,8 @@ abstract class FsHistoryProviderSuite extends SparkFunSuite with Matchers with P while (entry != null) { val actual = new String(ByteStreams.toByteArray(inputStream), StandardCharsets.UTF_8) val expected = - Files.toString(logs.find(_.getName == entry.getName).get, StandardCharsets.UTF_8) + Files.asCharSource(logs.find(_.getName == entry.getName).get, StandardCharsets.UTF_8) + .read() actual should be (expected) totalEntries += 1 entry = inputStream.getNextEntry diff --git a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerArgumentsSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerArgumentsSuite.scala index 2b9b110a41424..807e5ec3e823e 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerArgumentsSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerArgumentsSuite.scala @@ -45,8 +45,8 @@ class HistoryServerArgumentsSuite extends SparkFunSuite { test("Properties File Arguments Parsing --properties-file") { withTempDir { tmpDir => val outFile = File.createTempFile("test-load-spark-properties", "test", tmpDir) - Files.write("spark.test.CustomPropertyA blah\n" + - "spark.test.CustomPropertyB notblah\n", outFile, UTF_8) + Files.asCharSink(outFile, UTF_8).write("spark.test.CustomPropertyA blah\n" + + "spark.test.CustomPropertyB notblah\n") val argStrings = Array("--properties-file", outFile.getAbsolutePath) val hsa = new HistoryServerArguments(conf, argStrings) assert(conf.get("spark.test.CustomPropertyA") === "blah") diff --git a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala index abb5ae720af07..6b2bd90cd4314 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala @@ -283,7 +283,7 @@ abstract class HistoryServerSuite extends SparkFunSuite with BeforeAndAfter with val expectedFile = { new File(logDir, entry.getName) } - val expected = Files.toString(expectedFile, StandardCharsets.UTF_8) + val expected = Files.asCharSource(expectedFile, StandardCharsets.UTF_8).read() val actual = new String(ByteStreams.toByteArray(zipStream), StandardCharsets.UTF_8) actual should be (expected) filesCompared += 1 diff --git a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala index 79fa8d21bf3f1..fc8f48df2cb7d 100644 --- a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala +++ b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala @@ -383,7 +383,7 @@ object NonLocalModeSparkPlugin { resources: Map[String, ResourceInformation]): Unit = { val path = conf.get(TEST_PATH_CONF) val strToWrite = createFileStringWithGpuAddrs(id, resources) - Files.write(strToWrite, new File(path, s"$filePrefix$id"), StandardCharsets.UTF_8) + Files.asCharSink(new File(path, s"$filePrefix$id"), StandardCharsets.UTF_8).write(strToWrite) } def reset(): Unit = { diff --git a/core/src/test/scala/org/apache/spark/resource/ResourceDiscoveryPluginSuite.scala b/core/src/test/scala/org/apache/spark/resource/ResourceDiscoveryPluginSuite.scala index ff7d680352177..edf138df9e207 100644 --- a/core/src/test/scala/org/apache/spark/resource/ResourceDiscoveryPluginSuite.scala +++ b/core/src/test/scala/org/apache/spark/resource/ResourceDiscoveryPluginSuite.scala @@ -148,7 +148,7 @@ object TestResourceDiscoveryPlugin { def writeFile(conf: SparkConf, id: String): Unit = { val path = conf.get(TEST_PATH_CONF) val fileName = s"$id - ${UUID.randomUUID.toString}" - Files.write(id, new File(path, fileName), StandardCharsets.UTF_8) + Files.asCharSink(new File(path, fileName), StandardCharsets.UTF_8).write(id) } } diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala index 3ef382573517b..66b1ee7b58ac8 100644 --- a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala +++ b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala @@ -868,23 +868,23 @@ abstract class RpcEnvSuite extends SparkFunSuite { val conf = createSparkConf() val file = new File(tempDir, "file") - Files.write(UUID.randomUUID().toString(), file, UTF_8) + Files.asCharSink(file, UTF_8).write(UUID.randomUUID().toString) val fileWithSpecialChars = new File(tempDir, "file name") - Files.write(UUID.randomUUID().toString(), fileWithSpecialChars, UTF_8) + Files.asCharSink(fileWithSpecialChars, UTF_8).write(UUID.randomUUID().toString) val empty = new File(tempDir, "empty") - Files.write("", empty, UTF_8); + Files.asCharSink(empty, UTF_8).write("") val jar = new File(tempDir, "jar") - Files.write(UUID.randomUUID().toString(), jar, UTF_8) + Files.asCharSink(jar, UTF_8).write(UUID.randomUUID().toString) val dir1 = new File(tempDir, "dir1") assert(dir1.mkdir()) val subFile1 = new File(dir1, "file1") - Files.write(UUID.randomUUID().toString(), subFile1, UTF_8) + Files.asCharSink(subFile1, UTF_8).write(UUID.randomUUID().toString) val dir2 = new File(tempDir, "dir2") assert(dir2.mkdir()) val subFile2 = new File(dir2, "file2") - Files.write(UUID.randomUUID().toString(), subFile2, UTF_8) + Files.asCharSink(subFile2, UTF_8).write(UUID.randomUUID().toString) val fileUri = env.fileServer.addFile(file) val fileWithSpecialCharsUri = env.fileServer.addFile(fileWithSpecialChars) diff --git a/core/src/test/scala/org/apache/spark/util/FileAppenderSuite.scala b/core/src/test/scala/org/apache/spark/util/FileAppenderSuite.scala index 35ef0587b9b4c..4497ea1b2b798 100644 --- a/core/src/test/scala/org/apache/spark/util/FileAppenderSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/FileAppenderSuite.scala @@ -54,11 +54,11 @@ class FileAppenderSuite extends SparkFunSuite with BeforeAndAfter { val inputStream = new ByteArrayInputStream(testString.getBytes(StandardCharsets.UTF_8)) // The `header` should not be covered val header = "Add header" - Files.write(header, testFile, StandardCharsets.UTF_8) + Files.asCharSink(testFile, StandardCharsets.UTF_8).write(header) val appender = new FileAppender(inputStream, testFile) inputStream.close() appender.awaitTermination() - assert(Files.toString(testFile, StandardCharsets.UTF_8) === header + testString) + assert(Files.asCharSource(testFile, StandardCharsets.UTF_8).read() === header + testString) } test("SPARK-35027: basic file appender - close stream") { @@ -392,7 +392,7 @@ class FileAppenderSuite extends SparkFunSuite with BeforeAndAfter { IOUtils.closeQuietly(inputStream) } } else { - Files.toString(file, StandardCharsets.UTF_8) + Files.asCharSource(file, StandardCharsets.UTF_8).read() } }.mkString("") assert(allText === expectedText) diff --git a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala index a694e08def89c..a6e3345fc600c 100644 --- a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala @@ -735,8 +735,8 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties { withTempDir { tmpDir => val outFile = File.createTempFile("test-load-spark-properties", "test", tmpDir) System.setProperty("spark.test.fileNameLoadB", "2") - Files.write("spark.test.fileNameLoadA true\n" + - "spark.test.fileNameLoadB 1\n", outFile, UTF_8) + Files.asCharSink(outFile, UTF_8).write("spark.test.fileNameLoadA true\n" + + "spark.test.fileNameLoadB 1\n") val properties = Utils.getPropertiesFromFile(outFile.getAbsolutePath) properties .filter { case (k, v) => k.startsWith("spark.")} @@ -765,7 +765,7 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties { val innerSourceDir = Utils.createTempDir(root = sourceDir.getPath) val sourceFile = File.createTempFile("someprefix", "somesuffix", innerSourceDir) val targetDir = new File(tempDir, "target-dir") - Files.write("some text", sourceFile, UTF_8) + Files.asCharSink(sourceFile, UTF_8).write("some text") val path = if (Utils.isWindows) { diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java index 0c11c40cfe7ed..1052f47ea496e 100644 --- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java +++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.regex.Pattern; +import com.google.common.io.FileWriteMode; import scala.Tuple2; import com.google.common.io.Files; @@ -152,7 +153,8 @@ private static JavaStreamingContext createContext(String ip, System.out.println(output); System.out.println("Dropped " + droppedWordsCounter.value() + " word(s) totally"); System.out.println("Appending to " + outputFile.getAbsolutePath()); - Files.append(output + "\n", outputFile, Charset.defaultCharset()); + Files.asCharSink(outputFile, Charset.defaultCharset(), FileWriteMode.APPEND) + .write(output + "\n"); }); return ssc; diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/RecoverableNetworkWordCount.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/RecoverableNetworkWordCount.scala index 98539d6494231..1ec6ee4abd327 100644 --- a/examples/src/main/scala/org/apache/spark/examples/streaming/RecoverableNetworkWordCount.scala +++ b/examples/src/main/scala/org/apache/spark/examples/streaming/RecoverableNetworkWordCount.scala @@ -21,7 +21,7 @@ package org.apache.spark.examples.streaming import java.io.File import java.nio.charset.Charset -import com.google.common.io.Files +import com.google.common.io.{Files, FileWriteMode} import org.apache.spark.{SparkConf, SparkContext} import org.apache.spark.broadcast.Broadcast @@ -134,7 +134,8 @@ object RecoverableNetworkWordCount { println(output) println(s"Dropped ${droppedWordsCounter.value} word(s) totally") println(s"Appending to ${outputFile.getAbsolutePath}") - Files.append(output + "\n", outputFile, Charset.defaultCharset()) + Files.asCharSink(outputFile, Charset.defaultCharset(), FileWriteMode.APPEND) + .write(output + "\n") } ssc } diff --git a/mllib/src/test/java/org/apache/spark/ml/source/libsvm/JavaLibSVMRelationSuite.java b/mllib/src/test/java/org/apache/spark/ml/source/libsvm/JavaLibSVMRelationSuite.java index c3038fa9e1f8f..5f0d22ea2a8aa 100644 --- a/mllib/src/test/java/org/apache/spark/ml/source/libsvm/JavaLibSVMRelationSuite.java +++ b/mllib/src/test/java/org/apache/spark/ml/source/libsvm/JavaLibSVMRelationSuite.java @@ -50,7 +50,7 @@ public void setUp() throws IOException { tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "datasource"); File file = new File(tempDir, "part-00000"); String s = "1 1:1.0 3:2.0 5:3.0\n0\n0 2:4.0 4:5.0 6:6.0"; - Files.write(s, file, StandardCharsets.UTF_8); + Files.asCharSink(file, StandardCharsets.UTF_8).write(s); path = tempDir.toURI().toString(); } diff --git a/mllib/src/test/scala/org/apache/spark/ml/source/libsvm/LibSVMRelationSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/source/libsvm/LibSVMRelationSuite.scala index f2bb145614725..6a0d7b1237ee4 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/source/libsvm/LibSVMRelationSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/source/libsvm/LibSVMRelationSuite.scala @@ -65,9 +65,9 @@ class LibSVMRelationSuite val succ = new File(dir, "_SUCCESS") val file0 = new File(dir, "part-00000") val file1 = new File(dir, "part-00001") - Files.write("", succ, StandardCharsets.UTF_8) - Files.write(lines0, file0, StandardCharsets.UTF_8) - Files.write(lines1, file1, StandardCharsets.UTF_8) + Files.asCharSink(succ, StandardCharsets.UTF_8).write("") + Files.asCharSink(file0, StandardCharsets.UTF_8).write(lines0) + Files.asCharSink(file1, StandardCharsets.UTF_8).write(lines1) path = dir.getPath } diff --git a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala index a90c9c80d4959..1a02e26b9260c 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala @@ -93,7 +93,7 @@ class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { """.stripMargin val tempDir = Utils.createTempDir() val file = new File(tempDir.getPath, "part-00000") - Files.write(lines, file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(lines) val path = tempDir.toURI.toString val pointsWithNumFeatures = loadLibSVMFile(sc, path, 6).collect() @@ -126,7 +126,7 @@ class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { """.stripMargin val tempDir = Utils.createTempDir() val file = new File(tempDir.getPath, "part-00000") - Files.write(lines, file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(lines) val path = tempDir.toURI.toString intercept[SparkException] { @@ -143,7 +143,7 @@ class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { """.stripMargin val tempDir = Utils.createTempDir() val file = new File(tempDir.getPath, "part-00000") - Files.write(lines, file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(lines) val path = tempDir.toURI.toString intercept[SparkException] { diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/SparkKubernetesClientFactory.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/SparkKubernetesClientFactory.scala index 79f76e96474e3..2c28dc380046c 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/SparkKubernetesClientFactory.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/SparkKubernetesClientFactory.scala @@ -107,7 +107,7 @@ object SparkKubernetesClientFactory extends Logging { (token, configBuilder) => configBuilder.withOauthToken(token) }.withOption(oauthTokenFile) { (file, configBuilder) => - configBuilder.withOauthToken(Files.toString(file, Charsets.UTF_8)) + configBuilder.withOauthToken(Files.asCharSource(file, Charsets.UTF_8).read()) }.withOption(caCertFile) { (file, configBuilder) => configBuilder.withCaCertFile(file) }.withOption(clientKeyFile) { diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStep.scala index e266d0f904e46..d64378a65d66f 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStep.scala @@ -116,7 +116,7 @@ private[spark] class HadoopConfDriverFeatureStep(conf: KubernetesConf) override def getAdditionalKubernetesResources(): Seq[HasMetadata] = { if (confDir.isDefined) { val fileMap = confFiles.map { file => - (file.getName(), Files.toString(file, StandardCharsets.UTF_8)) + (file.getName(), Files.asCharSource(file, StandardCharsets.UTF_8).read()) }.toMap.asJava Seq(new ConfigMapBuilder() diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStep.scala index 82bda88892d04..89aefe47e46d1 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStep.scala @@ -229,7 +229,7 @@ private[spark] class KerberosConfDriverFeatureStep(kubernetesConf: KubernetesDri .endMetadata() .withImmutable(true) .addToData( - Map(file.getName() -> Files.toString(file, StandardCharsets.UTF_8)).asJava) + Map(file.getName() -> Files.asCharSource(file, StandardCharsets.UTF_8).read()).asJava) .build() } } ++ { diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/PodTemplateConfigMapStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/PodTemplateConfigMapStep.scala index cdc0112294113..f94dad2d15dc1 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/PodTemplateConfigMapStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/PodTemplateConfigMapStep.scala @@ -81,7 +81,7 @@ private[spark] class PodTemplateConfigMapStep(conf: KubernetesConf) val hadoopConf = SparkHadoopUtil.get.newConfiguration(conf.sparkConf) val uri = downloadFile(podTemplateFile, Utils.createTempDir(), conf.sparkConf, hadoopConf) val file = new java.net.URI(uri).getPath - val podTemplateString = Files.toString(new File(file), StandardCharsets.UTF_8) + val podTemplateString = Files.asCharSource(new File(file), StandardCharsets.UTF_8).read() Seq(new ConfigMapBuilder() .withNewMetadata() .withName(configmapName) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala index f1dd8b94f17ff..a72152a851c4f 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala @@ -128,7 +128,7 @@ class DriverKubernetesCredentialsFeatureStepSuite extends SparkFunSuite { private def writeCredentials(credentialsFileName: String, credentialsContents: String): File = { val credentialsFile = new File(credentialsTempDirectory, credentialsFileName) - Files.write(credentialsContents, credentialsFile, Charsets.UTF_8) + Files.asCharSink(credentialsFile, Charsets.UTF_8).write(credentialsContents) credentialsFile } } diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStepSuite.scala index 8f21b95236a9c..4310ac0220e5e 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfDriverFeatureStepSuite.scala @@ -48,7 +48,7 @@ class HadoopConfDriverFeatureStepSuite extends SparkFunSuite { val confFiles = Set("core-site.xml", "hdfs-site.xml") confFiles.foreach { f => - Files.write("some data", new File(confDir, f), UTF_8) + Files.asCharSink(new File(confDir, f), UTF_8).write("some data") } val sparkConf = new SparkConfWithEnv(Map(ENV_HADOOP_CONF_DIR -> confDir.getAbsolutePath())) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfExecutorFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfExecutorFeatureStepSuite.scala index a60227814eb13..04e20258d068f 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfExecutorFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/HadoopConfExecutorFeatureStepSuite.scala @@ -36,7 +36,7 @@ class HadoopConfExecutorFeatureStepSuite extends SparkFunSuite { val confFiles = Set("core-site.xml", "hdfs-site.xml") confFiles.foreach { f => - Files.write("some data", new File(confDir, f), UTF_8) + Files.asCharSink(new File(confDir, f), UTF_8).write("some data") } Seq( diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStepSuite.scala index 163d87643abd3..b172bdc06ddca 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/KerberosConfDriverFeatureStepSuite.scala @@ -55,7 +55,7 @@ class KerberosConfDriverFeatureStepSuite extends SparkFunSuite { test("create krb5.conf config map if local config provided") { val krbConf = File.createTempFile("krb5", ".conf", tmpDir) - Files.write("some data", krbConf, UTF_8) + Files.asCharSink(krbConf, UTF_8).write("some data") val sparkConf = new SparkConf(false) .set(KUBERNETES_KERBEROS_KRB5_FILE, krbConf.getAbsolutePath()) @@ -70,7 +70,7 @@ class KerberosConfDriverFeatureStepSuite extends SparkFunSuite { test("create keytab secret if client keytab file used") { val keytab = File.createTempFile("keytab", ".bin", tmpDir) - Files.write("some data", keytab, UTF_8) + Files.asCharSink(keytab, UTF_8).write("some data") val sparkConf = new SparkConf(false) .set(KEYTAB, keytab.getAbsolutePath()) diff --git a/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/DecommissionSuite.scala b/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/DecommissionSuite.scala index ae5f037c6b7d4..950079dcb5362 100644 --- a/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/DecommissionSuite.scala +++ b/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/DecommissionSuite.scala @@ -40,7 +40,7 @@ private[spark] trait DecommissionSuite { k8sSuite: KubernetesSuite => val logConfFilePath = s"${sparkHomeDir.toFile}/conf/log4j2.properties" try { - Files.write( + Files.asCharSink(new File(logConfFilePath), StandardCharsets.UTF_8).write( """rootLogger.level = info |rootLogger.appenderRef.stdout.ref = console |appender.console.type = Console @@ -51,9 +51,7 @@ private[spark] trait DecommissionSuite { k8sSuite: KubernetesSuite => | |logger.spark.name = org.apache.spark |logger.spark.level = debug - """.stripMargin, - new File(logConfFilePath), - StandardCharsets.UTF_8) + """.stripMargin) f() } finally { diff --git a/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/KubernetesSuite.scala b/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/KubernetesSuite.scala index 0b0b30e5e04fd..cf129677ad9c2 100644 --- a/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/KubernetesSuite.scala +++ b/resource-managers/kubernetes/integration-tests/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/KubernetesSuite.scala @@ -129,7 +129,7 @@ class KubernetesSuite extends SparkFunSuite val tagFile = new File(path) require(tagFile.isFile, s"No file found for image tag at ${tagFile.getAbsolutePath}.") - Files.toString(tagFile, Charsets.UTF_8).trim + Files.asCharSource(tagFile, Charsets.UTF_8).read().trim } .orElse(sys.props.get(CONFIG_KEY_IMAGE_TAG)) .getOrElse { diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/BaseYarnClusterSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/BaseYarnClusterSuite.scala index f0177541accc1..e0dfac62847ea 100644 --- a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/BaseYarnClusterSuite.scala +++ b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/BaseYarnClusterSuite.scala @@ -86,7 +86,7 @@ abstract class BaseYarnClusterSuite extends SparkFunSuite with Matchers { logConfDir.mkdir() val logConfFile = new File(logConfDir, "log4j2.properties") - Files.write(LOG4J_CONF, logConfFile, StandardCharsets.UTF_8) + Files.asCharSink(logConfFile, StandardCharsets.UTF_8).write(LOG4J_CONF) // Disable the disk utilization check to avoid the test hanging when people's disks are // getting full. @@ -232,11 +232,11 @@ abstract class BaseYarnClusterSuite extends SparkFunSuite with Matchers { // an error message val output = new Object() { override def toString: String = outFile - .map(Files.toString(_, StandardCharsets.UTF_8)) + .map(Files.asCharSource(_, StandardCharsets.UTF_8).read()) .getOrElse("(stdout/stderr was not captured)") } assert(finalState === SparkAppHandle.State.FINISHED, output) - val resultString = Files.toString(result, StandardCharsets.UTF_8) + val resultString = Files.asCharSource(result, StandardCharsets.UTF_8).read() assert(resultString === expected, output) } diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala index 806efd39800fb..92d9f2d62d1c1 100644 --- a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala +++ b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala @@ -141,7 +141,7 @@ class YarnClusterSuite extends BaseYarnClusterSuite { | | |""".stripMargin - Files.write(coreSite, new File(customConf, "core-site.xml"), StandardCharsets.UTF_8) + Files.asCharSink(new File(customConf, "core-site.xml"), StandardCharsets.UTF_8).write(coreSite) val result = File.createTempFile("result", null, tempDir) val finalState = runSpark(false, @@ -295,23 +295,22 @@ class YarnClusterSuite extends BaseYarnClusterSuite { test("running Spark in yarn-cluster mode displays driver log links") { val log4jConf = new File(tempDir, "log4j.properties") val logOutFile = new File(tempDir, "logs") - Files.write( + Files.asCharSink(log4jConf, StandardCharsets.UTF_8).write( s"""rootLogger.level = debug |rootLogger.appenderRef.file.ref = file |appender.file.type = File |appender.file.name = file |appender.file.fileName = $logOutFile |appender.file.layout.type = PatternLayout - |""".stripMargin, - log4jConf, StandardCharsets.UTF_8) + |""".stripMargin) // Since this test is trying to extract log output from the SparkSubmit process itself, // standard options to the Spark process don't take effect. Leverage the java-opts file which // will get picked up for the SparkSubmit process. val confDir = new File(tempDir, "conf") confDir.mkdir() val javaOptsFile = new File(confDir, "java-opts") - Files.write(s"-Dlog4j.configurationFile=file://$log4jConf\n", javaOptsFile, - StandardCharsets.UTF_8) + Files.asCharSink(javaOptsFile, StandardCharsets.UTF_8) + .write(s"-Dlog4j.configurationFile=file://$log4jConf\n") val result = File.createTempFile("result", null, tempDir) val finalState = runSpark(clientMode = false, @@ -320,7 +319,7 @@ class YarnClusterSuite extends BaseYarnClusterSuite { extraEnv = Map("SPARK_CONF_DIR" -> confDir.getAbsolutePath), extraConf = Map(CLIENT_INCLUDE_DRIVER_LOGS_LINK.key -> true.toString)) checkResult(finalState, result) - val logOutput = Files.toString(logOutFile, StandardCharsets.UTF_8) + val logOutput = Files.asCharSource(logOutFile, StandardCharsets.UTF_8).read() val logFilePattern = raw"""(?s).+\sDriver Logs \(\): https?://.+/(\?\S+)?\s.+""" logOutput should fullyMatch regex logFilePattern.replace("", "stdout") logOutput should fullyMatch regex logFilePattern.replace("", "stderr") @@ -374,7 +373,7 @@ class YarnClusterSuite extends BaseYarnClusterSuite { extraEnv: Map[String, String] = Map()): Unit = { assume(isPythonAvailable) val primaryPyFile = new File(tempDir, "test.py") - Files.write(TEST_PYFILE, primaryPyFile, StandardCharsets.UTF_8) + Files.asCharSink(primaryPyFile, StandardCharsets.UTF_8).write(TEST_PYFILE) // When running tests, let's not assume the user has built the assembly module, which also // creates the pyspark archive. Instead, let's use PYSPARK_ARCHIVES_PATH to point at the @@ -396,7 +395,7 @@ class YarnClusterSuite extends BaseYarnClusterSuite { subdir } val pyModule = new File(moduleDir, "mod1.py") - Files.write(TEST_PYMODULE, pyModule, StandardCharsets.UTF_8) + Files.asCharSink(pyModule, StandardCharsets.UTF_8).write(TEST_PYMODULE) val mod2Archive = TestUtils.createJarWithFiles(Map("mod2.py" -> TEST_PYMODULE), moduleDir) val pyFiles = Seq(pyModule.getAbsolutePath(), mod2Archive.getPath()).mkString(",") @@ -443,7 +442,7 @@ class YarnClusterSuite extends BaseYarnClusterSuite { def createEmptyIvySettingsFile: File = { val emptyIvySettings = File.createTempFile("ivy", ".xml") - Files.write("", emptyIvySettings, StandardCharsets.UTF_8) + Files.asCharSink(emptyIvySettings, StandardCharsets.UTF_8).write("") emptyIvySettings } @@ -555,7 +554,7 @@ private object YarnClusterDriverUseSparkHadoopUtilConf extends Logging with Matc } result = "success" } finally { - Files.write(result, status, StandardCharsets.UTF_8) + Files.asCharSink(status, StandardCharsets.UTF_8).write(result) sc.stop() } } @@ -658,7 +657,7 @@ private object YarnClusterDriver extends Logging with Matchers { assert(driverAttributes === expectationAttributes) } } finally { - Files.write(result, status, StandardCharsets.UTF_8) + Files.asCharSink(status, StandardCharsets.UTF_8).write(result) sc.stop() } } @@ -707,7 +706,7 @@ private object YarnClasspathTest extends Logging { case t: Throwable => error(s"loading test.resource to $resultPath", t) } finally { - Files.write(result, new File(resultPath), StandardCharsets.UTF_8) + Files.asCharSink(new File(resultPath), StandardCharsets.UTF_8).write(result) } } @@ -751,7 +750,7 @@ private object YarnAddJarTest extends Logging { result = "success" } } finally { - Files.write(result, new File(resultPath), StandardCharsets.UTF_8) + Files.asCharSink(new File(resultPath), StandardCharsets.UTF_8).write(result) sc.stop() } } @@ -796,7 +795,7 @@ private object ExecutorEnvTestApp { executorEnvs.get(k).contains(v) } - Files.write(result.toString, new File(status), StandardCharsets.UTF_8) + Files.asCharSink(new File(status), StandardCharsets.UTF_8).write(result.toString) sc.stop() } diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnShuffleIntegrationSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnShuffleIntegrationSuite.scala index f745265eddfd9..f8d69c0ae568e 100644 --- a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnShuffleIntegrationSuite.scala +++ b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnShuffleIntegrationSuite.scala @@ -181,7 +181,7 @@ private object YarnExternalShuffleDriver extends Logging with Matchers { if (execStateCopy != null) { FileUtils.deleteDirectory(execStateCopy) } - Files.write(result, status, StandardCharsets.UTF_8) + Files.asCharSink(status, StandardCharsets.UTF_8).write(result) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala index 275b35947182c..c90b1d3ca5978 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowConvertersSuite.scala @@ -1217,8 +1217,8 @@ class ArrowConvertersSuite extends SharedSparkSession { val tempFile1 = new File(tempDataPath, "testData2-ints-part1.json") val tempFile2 = new File(tempDataPath, "testData2-ints-part2.json") - Files.write(json1, tempFile1, StandardCharsets.UTF_8) - Files.write(json2, tempFile2, StandardCharsets.UTF_8) + Files.asCharSink(tempFile1, StandardCharsets.UTF_8).write(json1) + Files.asCharSink(tempFile2, StandardCharsets.UTF_8).write(json2) validateConversion(schema, arrowBatches(0), tempFile1) validateConversion(schema, arrowBatches(1), tempFile2) @@ -1501,7 +1501,7 @@ class ArrowConvertersSuite extends SharedSparkSession { // NOTE: coalesce to single partition because can only load 1 batch in validator val batchBytes = df.coalesce(1).toArrowBatchRdd.collect().head val tempFile = new File(tempDataPath, file) - Files.write(json, tempFile, StandardCharsets.UTF_8) + Files.asCharSink(tempFile, StandardCharsets.UTF_8).write(json) validateConversion(df.schema, batchBytes, tempFile, timeZoneId, errorOnDuplicatedFieldNames) } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala index 4575549005f33..f1f0befcb0d30 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala @@ -1222,7 +1222,7 @@ abstract class HiveThriftServer2TestBase extends SparkFunSuite with BeforeAndAft // overrides all other potential log4j configurations contained in other dependency jar files. val tempLog4jConf = Utils.createTempDir().getCanonicalPath - Files.write( + Files.asCharSink(new File(s"$tempLog4jConf/log4j2.properties"), StandardCharsets.UTF_8).write( """rootLogger.level = info |rootLogger.appenderRef.stdout.ref = console |appender.console.type = Console @@ -1230,9 +1230,7 @@ abstract class HiveThriftServer2TestBase extends SparkFunSuite with BeforeAndAft |appender.console.target = SYSTEM_ERR |appender.console.layout.type = PatternLayout |appender.console.layout.pattern = %d{HH:mm:ss.SSS} %p %c: %maxLen{%m}{512}%n%ex{8}%n - """.stripMargin, - new File(s"$tempLog4jConf/log4j2.properties"), - StandardCharsets.UTF_8) + """.stripMargin) tempLog4jConf } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/UISeleniumSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/UISeleniumSuite.scala index 2b2cbec41d643..8d4a9886a2b25 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/UISeleniumSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/UISeleniumSuite.scala @@ -75,7 +75,7 @@ class UISeleniumSuite // overrides all other potential log4j configurations contained in other dependency jar files. val tempLog4jConf = org.apache.spark.util.Utils.createTempDir().getCanonicalPath - Files.write( + Files.asCharSink(new File(s"$tempLog4jConf/log4j2.properties"), StandardCharsets.UTF_8).write( """rootLogger.level = info |rootLogger.appenderRef.file.ref = console |appender.console.type = Console @@ -83,9 +83,7 @@ class UISeleniumSuite |appender.console.target = SYSTEM_ERR |appender.console.layout.type = PatternLayout |appender.console.layout.pattern = %d{HH:mm:ss.SSS} %p %c: %maxLen{%m}{512}%n%ex{8}%n - """.stripMargin, - new File(s"$tempLog4jConf/log4j2.properties"), - StandardCharsets.UTF_8) + """.stripMargin) tempLog4jConf } diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala index 14051034a588e..1c45b02375b30 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala @@ -23,7 +23,7 @@ import java.nio.charset.StandardCharsets import java.sql.{Date, Timestamp} import java.util.{Locale, Set} -import com.google.common.io.Files +import com.google.common.io.{Files, FileWriteMode} import org.apache.hadoop.fs.{FileSystem, Path} import org.apache.spark.{SparkException, TestUtils} @@ -1947,10 +1947,10 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile for (i <- 1 to 3) { - Files.write(s"$i", new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8).write(s"$i") } for (i <- 5 to 7) { - Files.write(s"$i", new File(dirPath, s"part-s-0000$i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-s-0000$i"), StandardCharsets.UTF_8).write(s"$i") } withTable("load_t") { @@ -1971,7 +1971,7 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile for (i <- 1 to 3) { - Files.write(s"$i", new File(dirPath, s"part-r-0000 $i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-r-0000 $i"), StandardCharsets.UTF_8).write(s"$i") } withTable("load_t") { sql("CREATE TABLE load_t (a STRING) USING hive") @@ -1986,7 +1986,7 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile for (i <- 1 to 3) { - Files.write(s"$i", new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8).write(s"$i") } withTable("load_t") { sql("CREATE TABLE load_t (a STRING) USING hive") @@ -2010,7 +2010,7 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile for (i <- 1 to 3) { - Files.write(s"$i", new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8).write(s"$i") } withTable("load_t1") { sql("CREATE TABLE load_t1 (a STRING) USING hive") @@ -2025,7 +2025,7 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile for (i <- 1 to 3) { - Files.write(s"$i", new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8) + Files.asCharSink(new File(dirPath, s"part-r-0000$i"), StandardCharsets.UTF_8).write(s"$i") } withTable("load_t2") { sql("CREATE TABLE load_t2 (a STRING) USING hive") @@ -2039,7 +2039,8 @@ abstract class SQLQuerySuiteBase extends QueryTest with SQLTestUtils with TestHi withTempDir { dir => val path = dir.toURI.toString.stripSuffix("/") val dirPath = dir.getAbsoluteFile - Files.append("1", new File(dirPath, "part-r-000011"), StandardCharsets.UTF_8) + Files.asCharSink( + new File(dirPath, "part-r-000011"), StandardCharsets.UTF_8, FileWriteMode.APPEND).write("1") withTable("part_table") { withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { sql( diff --git a/streaming/src/test/java/test/org/apache/spark/streaming/JavaAPISuite.java b/streaming/src/test/java/test/org/apache/spark/streaming/JavaAPISuite.java index f8d961fa8dd8e..73c2e89f3729a 100644 --- a/streaming/src/test/java/test/org/apache/spark/streaming/JavaAPISuite.java +++ b/streaming/src/test/java/test/org/apache/spark/streaming/JavaAPISuite.java @@ -1641,7 +1641,7 @@ public void testRawSocketStream() { private static List> fileTestPrepare(File testDir) throws IOException { File existingFile = new File(testDir, "0"); - Files.write("0\n", existingFile, StandardCharsets.UTF_8); + Files.asCharSink(existingFile, StandardCharsets.UTF_8).write("0\n"); Assertions.assertTrue(existingFile.setLastModified(1000)); Assertions.assertEquals(1000, existingFile.lastModified()); return Arrays.asList(Arrays.asList("0")); diff --git a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala index 43b0835df7cbf..4aeb0e043a973 100644 --- a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala +++ b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala @@ -649,7 +649,7 @@ class CheckpointSuite extends TestSuiteBase with LocalStreamingContext with DStr */ def writeFile(i: Int, clock: Clock): Unit = { val file = new File(testDir, i.toString) - Files.write(s"$i\n", file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(s"$i\n") assert(file.setLastModified(clock.getTimeMillis())) // Check that the file's modification date is actually the value we wrote, since rounding or // truncation will break the test: diff --git a/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala index 66fd1ac7bb22e..64335a96045bf 100644 --- a/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala +++ b/streaming/src/test/scala/org/apache/spark/streaming/InputStreamsSuite.scala @@ -132,7 +132,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter { val batchDuration = Seconds(2) // Create a file that exists before the StreamingContext is created: val existingFile = new File(testDir, "0") - Files.write("0\n", existingFile, StandardCharsets.UTF_8) + Files.asCharSink(existingFile, StandardCharsets.UTF_8).write("0\n") assert(existingFile.setLastModified(10000) && existingFile.lastModified === 10000) // Set up the streaming context and input streams @@ -191,7 +191,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter { // Create a file that exists before the StreamingContext is created: val existingFile = new File(testDir, "0") - Files.write("0\n", existingFile, StandardCharsets.UTF_8) + Files.asCharSink(existingFile, StandardCharsets.UTF_8).write("0\n") assert(existingFile.setLastModified(10000) && existingFile.lastModified === 10000) val pathWithWildCard = testDir.toString + "/*/" @@ -215,7 +215,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter { def createFileAndAdvanceTime(data: Int, dir: File): Unit = { val file = new File(testSubDir1, data.toString) - Files.write(s"$data\n", file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(s"$data\n") assert(file.setLastModified(clock.getTimeMillis())) assert(file.lastModified === clock.getTimeMillis()) logInfo(s"Created file $file") @@ -478,7 +478,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter { val batchDuration = Seconds(2) // Create a file that exists before the StreamingContext is created: val existingFile = new File(testDir, "0") - Files.write("0\n", existingFile, StandardCharsets.UTF_8) + Files.asCharSink(existingFile, StandardCharsets.UTF_8).write("0\n") assert(existingFile.setLastModified(10000) && existingFile.lastModified === 10000) // Set up the streaming context and input streams @@ -502,7 +502,7 @@ class InputStreamsSuite extends TestSuiteBase with BeforeAndAfter { val input = Seq(1, 2, 3, 4, 5) input.foreach { i => val file = new File(testDir, i.toString) - Files.write(s"$i\n", file, StandardCharsets.UTF_8) + Files.asCharSink(file, StandardCharsets.UTF_8).write(s"$i\n") assert(file.setLastModified(clock.getTimeMillis())) assert(file.lastModified === clock.getTimeMillis()) logInfo("Created file " + file) diff --git a/streaming/src/test/scala/org/apache/spark/streaming/MasterFailureTest.scala b/streaming/src/test/scala/org/apache/spark/streaming/MasterFailureTest.scala index 771e65ed40b51..2dc43a231d9b8 100644 --- a/streaming/src/test/scala/org/apache/spark/streaming/MasterFailureTest.scala +++ b/streaming/src/test/scala/org/apache/spark/streaming/MasterFailureTest.scala @@ -375,7 +375,7 @@ class FileGeneratingThread(input: Seq[String], testDir: Path, interval: Long) val localFile = new File(localTestDir, (i + 1).toString) val hadoopFile = new Path(testDir, (i + 1).toString) val tempHadoopFile = new Path(testDir, ".tmp_" + (i + 1).toString) - Files.write(input(i) + "\n", localFile, StandardCharsets.UTF_8) + Files.asCharSink(localFile, StandardCharsets.UTF_8).write(input(i) + "\n") var tries = 0 var done = false while (!done && tries < maxTries) { From 97e9bb3ac4b66711ced640ea466eeea5da6d1fd2 Mon Sep 17 00:00:00 2001 From: Gideon P Date: Tue, 1 Oct 2024 15:09:35 +0200 Subject: [PATCH 184/189] [SPARK-48700][SQL] Mode expression for complex types (all collations) ### What changes were proposed in this pull request? Add support for complex types with subfields that are collated strings, for the mode operator. ### Why are the changes needed? Full support for collations as per SPARK-48700 ### Does this PR introduce _any_ user-facing change? Yes. ### How was this patch tested? Unit tests only, so far. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #47154 from GideonPotok/collationmodecomplex. Lead-authored-by: Gideon P Co-authored-by: Gideon Potok <31429832+GideonPotok@users.noreply.github.com> Signed-off-by: Max Gekk --- .../resources/error/error-conditions.json | 10 + .../catalyst/expressions/aggregate/Mode.scala | 85 ++++-- .../sql/CollationSQLExpressionsSuite.scala | 257 ++++++++++++------ 3 files changed, 250 insertions(+), 102 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index fcaf2b1d9d301..3786643125a9f 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -631,6 +631,11 @@ "Cannot process input data types for the expression: ." ], "subClass" : { + "BAD_INPUTS" : { + "message" : [ + "The input data types to must be valid, but found the input types ." + ] + }, "MISMATCHED_TYPES" : { "message" : [ "All input types must be the same except nullable, containsNull, valueContainsNull flags, but found the input types ." @@ -1011,6 +1016,11 @@ "The input of can't be type data." ] }, + "UNSUPPORTED_MODE_DATA_TYPE" : { + "message" : [ + "The does not support the data type, because there is a \"MAP\" type with keys and/or values that have collated sub-fields." + ] + }, "UNSUPPORTED_UDF_INPUT_TYPE" : { "message" : [ "UDFs do not support '' as an input data type." diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala index e254a670991a1..8998348f0571b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala @@ -17,14 +17,17 @@ package org.apache.spark.sql.catalyst.expressions.aggregate +import org.apache.spark.SparkIllegalArgumentException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{ExpressionBuilder, TypeCheckResult, UnresolvedWithinGroup} import org.apache.spark.sql.catalyst.expressions.{Ascending, Descending, Expression, ExpressionDescription, ImplicitCastInputTypes, SortOrder} +import org.apache.spark.sql.catalyst.expressions.Cast.toSQLExpr import org.apache.spark.sql.catalyst.trees.UnaryLike import org.apache.spark.sql.catalyst.types.PhysicalDataType -import org.apache.spark.sql.catalyst.util.{CollationFactory, GenericArrayData, UnsafeRowUtils} +import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory, GenericArrayData, UnsafeRowUtils} +import org.apache.spark.sql.errors.DataTypeErrors.{toSQLId, toSQLType} import org.apache.spark.sql.errors.QueryCompilationErrors -import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType, StringType} +import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType, MapType, StringType, StructField, StructType} import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.collection.OpenHashMap @@ -50,17 +53,20 @@ case class Mode( override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType) override def checkInputDataTypes(): TypeCheckResult = { - if (UnsafeRowUtils.isBinaryStable(child.dataType) || child.dataType.isInstanceOf[StringType]) { + // TODO: SPARK-49358: Mode expression for map type with collated fields + if (UnsafeRowUtils.isBinaryStable(child.dataType) || + !child.dataType.existsRecursively(f => f.isInstanceOf[MapType] && + !UnsafeRowUtils.isBinaryStable(f))) { /* * The Mode class uses collation awareness logic to handle string data. - * Complex types with collated fields are not yet supported. + * All complex types except MapType with collated fields are supported. */ - // TODO: SPARK-48700: Mode expression for complex types (all collations) super.checkInputDataTypes() } else { - TypeCheckResult.TypeCheckFailure("The input to the function 'mode' was" + - " a type of binary-unstable type that is " + - s"not currently supported by ${prettyName}.") + TypeCheckResult.DataTypeMismatch("UNSUPPORTED_MODE_DATA_TYPE", + messageParameters = + Map("child" -> toSQLType(child.dataType), + "mode" -> toSQLId(prettyName))) } } @@ -86,6 +92,54 @@ case class Mode( buffer } + private def getCollationAwareBuffer( + childDataType: DataType, + buffer: OpenHashMap[AnyRef, Long]): Iterable[(AnyRef, Long)] = { + def groupAndReduceBuffer(groupingFunction: AnyRef => _): Iterable[(AnyRef, Long)] = { + buffer.groupMapReduce(t => + groupingFunction(t._1))(x => x)((x, y) => (x._1, x._2 + y._2)).values + } + def determineBufferingFunction( + childDataType: DataType): Option[AnyRef => _] = { + childDataType match { + case _ if UnsafeRowUtils.isBinaryStable(child.dataType) => None + case _ => Some(collationAwareTransform(_, childDataType)) + } + } + determineBufferingFunction(childDataType).map(groupAndReduceBuffer).getOrElse(buffer) + } + + protected[sql] def collationAwareTransform(data: AnyRef, dataType: DataType): AnyRef = { + dataType match { + case _ if UnsafeRowUtils.isBinaryStable(dataType) => data + case st: StructType => + processStructTypeWithBuffer(data.asInstanceOf[InternalRow].toSeq(st).zip(st.fields)) + case at: ArrayType => processArrayTypeWithBuffer(at, data.asInstanceOf[ArrayData]) + case st: StringType => + CollationFactory.getCollationKey(data.asInstanceOf[UTF8String], st.collationId) + case _ => + throw new SparkIllegalArgumentException( + errorClass = "COMPLEX_EXPRESSION_UNSUPPORTED_INPUT.BAD_INPUTS", + messageParameters = Map( + "expression" -> toSQLExpr(this), + "functionName" -> toSQLType(prettyName), + "dataType" -> toSQLType(child.dataType)) + ) + } + } + + private def processStructTypeWithBuffer( + tuples: Seq[(Any, StructField)]): Seq[Any] = { + tuples.map(t => collationAwareTransform(t._1.asInstanceOf[AnyRef], t._2.dataType)) + } + + private def processArrayTypeWithBuffer( + a: ArrayType, + data: ArrayData): Seq[Any] = { + (0 until data.numElements()).map(i => + collationAwareTransform(data.get(i, a.elementType), a.elementType)) + } + override def eval(buffer: OpenHashMap[AnyRef, Long]): Any = { if (buffer.isEmpty) { return null @@ -102,17 +156,12 @@ case class Mode( * to a single value (the sum of the counts), and finally reduces the groups to a single map. * * The new map is then used in the rest of the Mode evaluation logic. + * + * It is expected to work for all simple and complex types with + * collated fields, except for MapType (temporarily). */ - val collationAwareBuffer = child.dataType match { - case c: StringType if - !CollationFactory.fetchCollation(c.collationId).supportsBinaryEquality => - val collationId = c.collationId - val modeMap = buffer.toSeq.groupMapReduce { - case (k, _) => CollationFactory.getCollationKey(k.asInstanceOf[UTF8String], collationId) - }(x => x)((x, y) => (x._1, x._2 + y._2)).values - modeMap - case _ => buffer - } + val collationAwareBuffer = getCollationAwareBuffer(child.dataType, buffer) + reverseOpt.map { reverse => val defaultKeyOrdering = if (reverse) { PhysicalDataType.ordering(child.dataType).asInstanceOf[Ordering[AnyRef]].reverse diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala index 941d5cd31db40..9930709cd8bf3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala @@ -19,11 +19,12 @@ package org.apache.spark.sql import java.sql.{Date, Timestamp} import java.text.SimpleDateFormat +import java.util.Locale import scala.collection.immutable.Seq -import org.apache.spark.{SparkConf, SparkException, SparkIllegalArgumentException, SparkRuntimeException} -import org.apache.spark.sql.catalyst.ExtendedAnalysisException +import org.apache.spark.{SparkConf, SparkException, SparkIllegalArgumentException, SparkRuntimeException, SparkThrowable} +import org.apache.spark.sql.catalyst.{ExtendedAnalysisException, InternalRow} import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate.Mode import org.apache.spark.sql.internal.{SqlApiConf, SQLConf} @@ -1752,7 +1753,7 @@ class CollationSQLExpressionsSuite UTF8StringModeTestCase("unicode_ci", bufferValuesUTF8String, "b"), UTF8StringModeTestCase("unicode", bufferValuesUTF8String, "a")) - testCasesUTF8String.foreach(t => { + testCasesUTF8String.foreach ( t => { val buffer = new OpenHashMap[AnyRef, Long](5) val myMode = Mode(child = Literal.create("some_column_name", StringType(t.collationId))) t.bufferValues.foreach { case (k, v) => buffer.update(k, v) } @@ -1760,6 +1761,40 @@ class CollationSQLExpressionsSuite }) } + test("Support Mode.eval(buffer) with complex types") { + case class UTF8StringModeTestCase[R]( + collationId: String, + bufferValues: Map[InternalRow, Long], + result: R) + + val bufferValuesUTF8String: Map[Any, Long] = Map( + UTF8String.fromString("a") -> 5L, + UTF8String.fromString("b") -> 4L, + UTF8String.fromString("B") -> 3L, + UTF8String.fromString("d") -> 2L, + UTF8String.fromString("e") -> 1L) + + val bufferValuesComplex = bufferValuesUTF8String.map{ + case (k, v) => (InternalRow.fromSeq(Seq(k, k, k)), v) + } + val testCasesUTF8String = Seq( + UTF8StringModeTestCase("utf8_binary", bufferValuesComplex, "[a,a,a]"), + UTF8StringModeTestCase("UTF8_LCASE", bufferValuesComplex, "[b,b,b]"), + UTF8StringModeTestCase("unicode_ci", bufferValuesComplex, "[b,b,b]"), + UTF8StringModeTestCase("unicode", bufferValuesComplex, "[a,a,a]")) + + testCasesUTF8String.foreach { t => + val buffer = new OpenHashMap[AnyRef, Long](5) + val myMode = Mode(child = Literal.create(null, StructType(Seq( + StructField("f1", StringType(t.collationId), true), + StructField("f2", StringType(t.collationId), true), + StructField("f3", StringType(t.collationId), true) + )))) + t.bufferValues.foreach { case (k, v) => buffer.update(k, v) } + assert(myMode.eval(buffer).toString.toLowerCase() == t.result.toLowerCase()) + } + } + test("Support mode for string expression with collated strings in struct") { case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R) val testCases = Seq( @@ -1780,33 +1815,7 @@ class CollationSQLExpressionsSuite t.collationId + ", f2: INT>) USING parquet") sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd) val query = s"SELECT lower(mode(i).f1) FROM ${tableName}" - if(t.collationId == "UTF8_LCASE" || - t.collationId == "unicode_ci" || - t.collationId == "unicode") { - // Cannot resolve "mode(i)" due to data type mismatch: - // Input to function mode was a complex type with strings collated on non-binary - // collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13; - val params = Seq(("sqlExpr", "\"mode(i)\""), - ("msg", "The input to the function 'mode'" + - " was a type of binary-unstable type that is not currently supported by mode."), - ("hint", "")).toMap - checkError( - exception = intercept[AnalysisException] { - sql(query) - }, - condition = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT", - parameters = params, - queryContext = Array( - ExpectedContext(objectType = "", - objectName = "", - startIndex = 13, - stopIndex = 19, - fragment = "mode(i)") - ) - ) - } else { - checkAnswer(sql(query), Row(t.result)) - } + checkAnswer(sql(query), Row(t.result)) } }) } @@ -1819,47 +1828,21 @@ class CollationSQLExpressionsSuite ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b") ) - testCases.foreach(t => { + testCases.foreach { t => val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) => (0L to numRepeats).map(_ => s"named_struct('f1', " + s"named_struct('f2', collate('$elt', '${t.collationId}')), 'f3', 1)").mkString(",") }.mkString(",") - val tableName = s"t_${t.collationId}_mode_nested_struct" + val tableName = s"t_${t.collationId}_mode_nested_struct1" withTable(tableName) { sql(s"CREATE TABLE ${tableName}(i STRUCT, f3: INT>) USING parquet") sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd) val query = s"SELECT lower(mode(i).f1.f2) FROM ${tableName}" - if(t.collationId == "UTF8_LCASE" || - t.collationId == "unicode_ci" || - t.collationId == "unicode") { - // Cannot resolve "mode(i)" due to data type mismatch: - // Input to function mode was a complex type with strings collated on non-binary - // collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13; - val params = Seq(("sqlExpr", "\"mode(i)\""), - ("msg", "The input to the function 'mode' " + - "was a type of binary-unstable type that is not currently supported by mode."), - ("hint", "")).toMap - checkError( - exception = intercept[AnalysisException] { - sql(query) - }, - condition = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT", - parameters = params, - queryContext = Array( - ExpectedContext(objectType = "", - objectName = "", - startIndex = 13, - stopIndex = 19, - fragment = "mode(i)") - ) - ) - } else { - checkAnswer(sql(query), Row(t.result)) - } + checkAnswer(sql(query), Row(t.result)) } - }) + } } test("Support mode for string expression with collated strings in array complex type") { @@ -1870,44 +1853,150 @@ class CollationSQLExpressionsSuite ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b") ) - testCases.foreach(t => { + testCases.foreach { t => + val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) => + (0L to numRepeats).map(_ => s"array(named_struct('f2', " + + s"collate('$elt', '${t.collationId}'), 'f3', 1))").mkString(",") + }.mkString(",") + + val tableName = s"t_${t.collationId}_mode_nested_struct2" + withTable(tableName) { + sql(s"CREATE TABLE ${tableName}(" + + s"i ARRAY< STRUCT>)" + + s" USING parquet") + sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd) + val query = s"SELECT lower(element_at(mode(i).f2, 1)) FROM ${tableName}" + checkAnswer(sql(query), Row(t.result)) + } + } + } + + test("Support mode for string expression with collated strings in 3D array type") { + case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R) + val testCases = Seq( + ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), + ModeTestCase("UTF8_LCASE", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"), + ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), + ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b") + ) + testCases.foreach { t => + val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) => + (0L to numRepeats).map(_ => + s"array(array(array(collate('$elt', '${t.collationId}'))))").mkString(",") + }.mkString(",") + + val tableName = s"t_${t.collationId}_mode_nested_3d_array" + withTable(tableName) { + sql(s"CREATE TABLE ${tableName}(i ARRAY>>) USING parquet") + sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd) + val query = s"SELECT lower(" + + s"element_at(element_at(element_at(mode(i),1),1),1)) FROM ${tableName}" + checkAnswer(sql(query), Row(t.result)) + } + } + } + + test("Support mode for string expression with collated complex type - Highly nested") { + case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R) + val testCases = Seq( + ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), + ModeTestCase("UTF8_LCASE", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"), + ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), + ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b") + ) + testCases.foreach { t => val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) => (0L to numRepeats).map(_ => s"array(named_struct('s1', named_struct('a2', " + s"array(collate('$elt', '${t.collationId}'))), 'f3', 1))").mkString(",") }.mkString(",") - val tableName = s"t_${t.collationId}_mode_nested_struct" + val tableName = s"t_${t.collationId}_mode_highly_nested_struct" withTable(tableName) { sql(s"CREATE TABLE ${tableName}(" + s"i ARRAY>, f3: INT>>)" + s" USING parquet") sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd) val query = s"SELECT lower(element_at(element_at(mode(i), 1).s1.a2, 1)) FROM ${tableName}" - if(t.collationId == "UTF8_LCASE" || - t.collationId == "unicode_ci" || t.collationId == "unicode") { - val params = Seq(("sqlExpr", "\"mode(i)\""), - ("msg", "The input to the function 'mode' was a type" + - " of binary-unstable type that is not currently supported by mode."), - ("hint", "")).toMap - checkError( - exception = intercept[AnalysisException] { - sql(query) - }, - condition = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT", - parameters = params, - queryContext = Array( - ExpectedContext(objectType = "", - objectName = "", - startIndex = 35, - stopIndex = 41, - fragment = "mode(i)") - ) - ) - } else { + checkAnswer(sql(query), Row(t.result)) + } + } + } + + test("Support mode expression with collated in recursively nested struct with map with keys") { + case class ModeTestCase(collationId: String, bufferValues: Map[String, Long], result: String) + Seq( + ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "{a -> 1}"), + ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "{a -> 1}"), + ModeTestCase("utf8_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "{b -> 1}"), + ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "{b -> 1}") + ).foreach { t1 => + def checkThisError(t: ModeTestCase, query: String): Any = { + val c = s"STRUCT>" + val c1 = s"\"${c}\"" + checkError( + exception = intercept[SparkThrowable] { + sql(query).collect() + }, + condition = "DATATYPE_MISMATCH.UNSUPPORTED_MODE_DATA_TYPE", + parameters = Map( + ("sqlExpr", "\"mode(i)\""), + ("child", c1), + ("mode", "`mode`")), + queryContext = Seq(ExpectedContext("mode(i)", 18, 24)).toArray + ) + } + + def getValuesToAdd(t: ModeTestCase): String = { + val valuesToAdd = t.bufferValues.map { + case (elt, numRepeats) => + (0L to numRepeats).map(i => + s"named_struct('m1', map(collate('$elt', '${t.collationId}'), 1))" + ).mkString(",") + }.mkString(",") + valuesToAdd + } + val tableName = s"t_${t1.collationId}_mode_nested_map_struct1" + withTable(tableName) { + sql(s"CREATE TABLE ${tableName}(" + + s"i STRUCT>) USING parquet") + sql(s"INSERT INTO ${tableName} VALUES ${getValuesToAdd(t1)}") + val query = "SELECT lower(cast(mode(i).m1 as string))" + + s" FROM ${tableName}" + if (t1.collationId == "utf8_binary") { + checkAnswer(sql(query), Row(t1.result)) + } else { + checkThisError(t1, query) } } - }) + } + } + + test("UDT with collation - Mode (throw exception)") { + case class ModeTestCase(collationId: String, bufferValues: Map[String, Long], result: String) + Seq( + ModeTestCase("utf8_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"), + ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"), + ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b") + ).foreach { t1 => + checkError( + exception = intercept[SparkIllegalArgumentException] { + Mode( + child = Literal.create(null, + MapType(StringType(t1.collationId), IntegerType)) + ).collationAwareTransform( + data = Map.empty[String, Any], + dataType = MapType(StringType(t1.collationId), IntegerType) + ) + }, + condition = "COMPLEX_EXPRESSION_UNSUPPORTED_INPUT.BAD_INPUTS", + parameters = Map( + "expression" -> "\"mode(NULL)\"", + "functionName" -> "\"MODE\"", + "dataType" -> s"\"MAP\"") + ) + } } test("SPARK-48430: Map value extraction with collations") { From 3093ad68d2a3c6bab9c1605381d27e700766be22 Mon Sep 17 00:00:00 2001 From: exmy Date: Tue, 1 Oct 2024 15:22:29 +0200 Subject: [PATCH 185/189] [MINOR] Fix a typo in First aggregate expression ### What changes were proposed in this pull request? Find a typo for the comment on code `mergeExpressions` of `First` aggregate expression, fix from `first.right` to `first.left`. ### Why are the changes needed? Fix typo, it's confused. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? N.A ### Was this patch authored or co-authored using generative AI tooling? No Closes #48298 from exmy/fix-comment. Authored-by: exmy Signed-off-by: Max Gekk --- .../apache/spark/sql/catalyst/expressions/aggregate/First.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala index 4fe00099ddc91..9a39a6fe98796 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala @@ -104,7 +104,7 @@ case class First(child: Expression, ignoreNulls: Boolean) override lazy val mergeExpressions: Seq[Expression] = { // For first, we can just check if valueSet.left is set to true. If it is set - // to true, we use first.right. If not, we use first.right (even if valueSet.right is + // to true, we use first.left. If not, we use first.right (even if valueSet.right is // false, we are safe to do so because first.right will be null in this case). Seq( /* first = */ If(valueSet.left, first.left, first.right), From 3551a9ee6d388f68f326cce1c0c9dad51e33ef58 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Tue, 1 Oct 2024 21:14:24 -0700 Subject: [PATCH 186/189] [SPARK-49845][CORE] Make `appArgs` and `environmentVariables` optional in REST API ### What changes were proposed in this pull request? This PR aims to make `appArgs` and `environmentVariables` fields optional in REST API. ### Why are the changes needed? `appArgs` and `environmentVariables` became mandatory due to the Apache Mesos limitation at Spark 2.2.2. Technically, this is a revert of SPARK-22574. - https://github.com/apache/spark/pull/19966 Since Apache Spark 4.0 removed Mesos support, we don't need these requirements. - https://github.com/apache/spark/pull/43135 ### Does this PR introduce _any_ user-facing change? No because this is a relaxation of enforcement. ### How was this patch tested? Pass the CIs with the revised test case. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48316 from dongjoon-hyun/SPARK-49845. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../org/apache/spark/deploy/rest/StandaloneRestServer.scala | 5 +++-- .../apache/spark/deploy/rest/SubmitRestProtocolRequest.scala | 2 -- .../apache/spark/deploy/rest/SubmitRestProtocolSuite.scala | 2 -- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/deploy/rest/StandaloneRestServer.scala b/core/src/main/scala/org/apache/spark/deploy/rest/StandaloneRestServer.scala index 31673f666173a..c92e79381ca9b 100644 --- a/core/src/main/scala/org/apache/spark/deploy/rest/StandaloneRestServer.scala +++ b/core/src/main/scala/org/apache/spark/deploy/rest/StandaloneRestServer.scala @@ -218,11 +218,12 @@ private[rest] class StandaloneSubmitRequestServlet( val (_, masterPort) = Utils.extractHostPortFromSparkUrl(masterUrl) val updatedMasters = masters.map( _.replace(s":$masterRestPort", s":$masterPort")).getOrElse(masterUrl) - val appArgs = request.appArgs + val appArgs = Option(request.appArgs).getOrElse(Array[String]()) // Filter SPARK_LOCAL_(IP|HOSTNAME) environment variables from being set on the remote system. // In addition, the placeholders are replaced into the values of environment variables. val environmentVariables = - request.environmentVariables.filterNot(x => x._1.matches("SPARK_LOCAL_(IP|HOSTNAME)")) + Option(request.environmentVariables).getOrElse(Map.empty[String, String]) + .filterNot(x => x._1.matches("SPARK_LOCAL_(IP|HOSTNAME)")) .map(x => (x._1, replacePlaceHolder(x._2))) // Construct driver description diff --git a/core/src/main/scala/org/apache/spark/deploy/rest/SubmitRestProtocolRequest.scala b/core/src/main/scala/org/apache/spark/deploy/rest/SubmitRestProtocolRequest.scala index 7f462148c71a1..63882259adcb5 100644 --- a/core/src/main/scala/org/apache/spark/deploy/rest/SubmitRestProtocolRequest.scala +++ b/core/src/main/scala/org/apache/spark/deploy/rest/SubmitRestProtocolRequest.scala @@ -47,8 +47,6 @@ private[rest] class CreateSubmissionRequest extends SubmitRestProtocolRequest { super.doValidate() assert(sparkProperties != null, "No Spark properties set!") assertFieldIsSet(appResource, "appResource") - assertFieldIsSet(appArgs, "appArgs") - assertFieldIsSet(environmentVariables, "environmentVariables") assertPropertyIsSet("spark.app.name") assertPropertyIsBoolean(config.DRIVER_SUPERVISE.key) assertPropertyIsNumeric(config.DRIVER_CORES.key) diff --git a/core/src/test/scala/org/apache/spark/deploy/rest/SubmitRestProtocolSuite.scala b/core/src/test/scala/org/apache/spark/deploy/rest/SubmitRestProtocolSuite.scala index 9eb5172583120..f2807f258f2d1 100644 --- a/core/src/test/scala/org/apache/spark/deploy/rest/SubmitRestProtocolSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/rest/SubmitRestProtocolSuite.scala @@ -87,8 +87,6 @@ class SubmitRestProtocolSuite extends SparkFunSuite { message.clientSparkVersion = "1.2.3" message.appResource = "honey-walnut-cherry.jar" message.mainClass = "org.apache.spark.examples.SparkPie" - message.appArgs = Array("two slices") - message.environmentVariables = Map("PATH" -> "/dev/null") val conf = new SparkConf(false) conf.set("spark.app.name", "SparkPie") message.sparkProperties = conf.getAll.toMap From 077a31989c99cb6302a325c953d2ee92ba573a8b Mon Sep 17 00:00:00 2001 From: Stefan Kandic Date: Wed, 2 Oct 2024 14:46:46 +0200 Subject: [PATCH 187/189] [SPARK-49843][SQL] Fix change comment on char/varchar columns ### What changes were proposed in this pull request? Fix the issue in `AlterTableChangeColumnCommand` where changing the comment of a char/varchar column also tries to change the column type to string. ### Why are the changes needed? Because the newColumn will always be a `StringType` even when the metadata says that it was originally char/varchar. ### Does this PR introduce _any_ user-facing change? Yes, the query will no longer fail when using this code path. ### How was this patch tested? New query in golden files. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48315 from stefankandic/fixAlterVarcharCol. Authored-by: Stefan Kandic Signed-off-by: Max Gekk --- .../analysis/ResolveSessionCatalog.scala | 10 ++++-- .../analyzer-results/charvarchar.sql.out | 12 +++++++ .../sql-tests/inputs/charvarchar.sql | 2 ++ .../sql-tests/results/charvarchar.sql.out | 32 ++++++++++++++----- 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala index a9ad7523c8fbc..884c870e8eed3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute} import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, toPrettySQL, ResolveDefaultColumns => DefaultCols} +import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, toPrettySQL, CharVarcharUtils, ResolveDefaultColumns => DefaultCols} import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns._ import org.apache.spark.sql.connector.catalog.{CatalogExtension, CatalogManager, CatalogPlugin, CatalogV2Util, LookupCatalog, SupportsNamespaces, V1Table} import org.apache.spark.sql.connector.expressions.Transform @@ -36,7 +36,7 @@ import org.apache.spark.sql.execution.datasources.{CreateTable => CreateTableV1} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Utils import org.apache.spark.sql.internal.{HiveSerDe, SQLConf} import org.apache.spark.sql.internal.connector.V1Function -import org.apache.spark.sql.types.{MetadataBuilder, StructField, StructType} +import org.apache.spark.sql.types.{MetadataBuilder, StringType, StructField, StructType} import org.apache.spark.util.ArrayImplicits._ /** @@ -87,7 +87,11 @@ class ResolveSessionCatalog(val catalogManager: CatalogManager) val colName = a.column.name(0) val dataType = a.dataType.getOrElse { table.schema.findNestedField(Seq(colName), resolver = conf.resolver) - .map(_._2.dataType) + .map { + case (_, StructField(_, st: StringType, _, metadata)) => + CharVarcharUtils.getRawType(metadata).getOrElse(st) + case (_, field) => field.dataType + } .getOrElse { throw QueryCompilationErrors.unresolvedColumnError( toSQLId(a.column.name), table.schema.fieldNames) diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar.sql.out index 5c1417f7c0aae..524797015a2f6 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar.sql.out @@ -263,6 +263,18 @@ desc formatted char_part DescribeTableCommand `spark_catalog`.`default`.`char_part`, true, [col_name#x, data_type#x, comment#x] +-- !query +alter table char_part change column c1 comment 'char comment' +-- !query analysis +AlterTableChangeColumnCommand `spark_catalog`.`default`.`char_part`, c1, StructField(c1,CharType(5),true) + + +-- !query +alter table char_part change column v1 comment 'varchar comment' +-- !query analysis +AlterTableChangeColumnCommand `spark_catalog`.`default`.`char_part`, v1, StructField(v1,VarcharType(6),true) + + -- !query alter table char_part add partition (v2='ke', c2='nt') location 'loc1' -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/charvarchar.sql b/sql/core/src/test/resources/sql-tests/inputs/charvarchar.sql index 8117dec53f4ab..be038e1083cd8 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/charvarchar.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/charvarchar.sql @@ -49,6 +49,8 @@ desc formatted char_tbl1; create table char_part(c1 char(5), c2 char(2), v1 varchar(6), v2 varchar(2)) using parquet partitioned by (v2, c2); desc formatted char_part; +alter table char_part change column c1 comment 'char comment'; +alter table char_part change column v1 comment 'varchar comment'; alter table char_part add partition (v2='ke', c2='nt') location 'loc1'; desc formatted char_part; diff --git a/sql/core/src/test/resources/sql-tests/results/charvarchar.sql.out b/sql/core/src/test/resources/sql-tests/results/charvarchar.sql.out index 568c9f3b29e87..8aafa25c5caaf 100644 --- a/sql/core/src/test/resources/sql-tests/results/charvarchar.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/charvarchar.sql.out @@ -556,6 +556,22 @@ Location [not included in comparison]/{warehouse_dir}/char_part Partition Provider Catalog +-- !query +alter table char_part change column c1 comment 'char comment' +-- !query schema +struct<> +-- !query output + + + +-- !query +alter table char_part change column v1 comment 'varchar comment' +-- !query schema +struct<> +-- !query output + + + -- !query alter table char_part add partition (v2='ke', c2='nt') location 'loc1' -- !query schema @@ -569,8 +585,8 @@ desc formatted char_part -- !query schema struct -- !query output -c1 char(5) -v1 varchar(6) +c1 char(5) char comment +v1 varchar(6) varchar comment v2 varchar(2) c2 char(2) # Partition Information @@ -612,8 +628,8 @@ desc formatted char_part -- !query schema struct -- !query output -c1 char(5) -v1 varchar(6) +c1 char(5) char comment +v1 varchar(6) varchar comment v2 varchar(2) c2 char(2) # Partition Information @@ -647,8 +663,8 @@ desc formatted char_part -- !query schema struct -- !query output -c1 char(5) -v1 varchar(6) +c1 char(5) char comment +v1 varchar(6) varchar comment v2 varchar(2) c2 char(2) # Partition Information @@ -682,8 +698,8 @@ desc formatted char_part -- !query schema struct -- !query output -c1 char(5) -v1 varchar(6) +c1 char(5) char comment +v1 varchar(6) varchar comment v2 varchar(2) c2 char(2) # Partition Information From 18dbaa5a070c74007137780e8529321b75b10b48 Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Wed, 2 Oct 2024 13:19:03 -0700 Subject: [PATCH 188/189] [SPARK-49560][SQL] Add SQL pipe syntax for the TABLESAMPLE operator ### What changes were proposed in this pull request? WIP This PR adds SQL pipe syntax support for the TABLESAMPLE operator. For example: ``` CREATE TABLE t(x INT, y STRING) USING CSV; INSERT INTO t VALUES (0, 'abc'), (1, 'def'); TABLE t |> TABLESAMPLE (100 PERCENT) REPEATABLE (0) |> TABLESAMPLE (5 ROWS) REPEATABLE (0) |> TABLESAMPLE (BUCKET 1 OUT OF 1) REPEATABLE (0); 0 abc 1 def ``` ### Why are the changes needed? The SQL pipe operator syntax will let users compose queries in a more flexible fashion. ### Does this PR introduce _any_ user-facing change? Yes, see above. ### How was this patch tested? This PR adds a few unit test cases, but mostly relies on golden file test coverage. I did this to make sure the answers are correct as this feature is implemented and also so we can look at the analyzer output plans to ensure they look right as well. ### Was this patch authored or co-authored using generative AI tooling? No Closes #48168 from dtenedor/pipe-tablesample. Authored-by: Daniel Tenedorio Signed-off-by: Gengliang Wang --- .../sql/catalyst/parser/SqlBaseParser.g4 | 1 + .../sql/catalyst/parser/AstBuilder.scala | 4 +- .../analyzer-results/pipe-operators.sql.out | 184 ++++++++++++++++ .../sql-tests/inputs/pipe-operators.sql | 49 +++++ .../sql-tests/results/pipe-operators.sql.out | 198 ++++++++++++++++++ .../sql/execution/SparkSqlParserSuite.scala | 9 + 6 files changed, 444 insertions(+), 1 deletion(-) diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index 33ac3249eb663..e8e2e980135a2 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -1504,6 +1504,7 @@ operatorPipeRightSide // messages in the event that both are present (this is not allowed). | pivotClause unpivotClause? | unpivotClause pivotClause? + | sample ; // When `SQL_standard_keyword_behavior=true`, there are 2 kinds of keywords in Spark SQL. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index e2350474a8708..9ce96ae652fed 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -5903,7 +5903,9 @@ class AstBuilder extends DataTypeAstBuilder throw QueryParsingErrors.unpivotWithPivotInFromClauseNotAllowedError(ctx) } withUnpivot(c, left) - }.get))) + }.getOrElse(Option(ctx.sample).map { c => + withSample(c, left) + }.get)))) } /** diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out index 8cd062aeb01a3..aee8da46aafbe 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/pipe-operators.sql.out @@ -921,6 +921,190 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +table t +|> tablesample (100 percent) repeatable (0) +-- !query analysis +Sample 0.0, 1.0, false, 0 ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> tablesample (2 rows) repeatable (0) +-- !query analysis +GlobalLimit 2 ++- LocalLimit 2 + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> tablesample (bucket 1 out of 1) repeatable (0) +-- !query analysis +Sample 0.0, 1.0, false, 0 ++- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> tablesample (100 percent) repeatable (0) +|> tablesample (5 rows) repeatable (0) +|> tablesample (bucket 1 out of 1) repeatable (0) +-- !query analysis +Sample 0.0, 1.0, false, 0 ++- GlobalLimit 5 + +- LocalLimit 5 + +- Sample 0.0, 1.0, false, 0 + +- SubqueryAlias spark_catalog.default.t + +- Relation spark_catalog.default.t[x#x,y#x] csv + + +-- !query +table t +|> tablesample () +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0014", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 25, + "fragment" : "tablesample ()" + } ] +} + + +-- !query +table t +|> tablesample (-100 percent) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0064", + "messageParameters" : { + "msg" : "Sampling fraction (-1.0) must be on interval [0, 1]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 37, + "fragment" : "tablesample (-100 percent)" + } ] +} + + +-- !query +table t +|> tablesample (-5 rows) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_LIMIT_LIKE_EXPRESSION.IS_NEGATIVE", + "sqlState" : "42K0E", + "messageParameters" : { + "expr" : "\"-5\"", + "name" : "limit", + "v" : "-5" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 25, + "stopIndex" : 26, + "fragment" : "-5" + } ] +} + + +-- !query +table t +|> tablesample (x rows) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_LIMIT_LIKE_EXPRESSION.IS_UNFOLDABLE", + "sqlState" : "42K0E", + "messageParameters" : { + "expr" : "\"x\"", + "name" : "limit" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 25, + "stopIndex" : 25, + "fragment" : "x" + } ] +} + + +-- !query +table t +|> tablesample (bucket 2 out of 1) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0064", + "messageParameters" : { + "msg" : "Sampling fraction (2.0) must be on interval [0, 1]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 42, + "fragment" : "tablesample (bucket 2 out of 1)" + } ] +} + + +-- !query +table t +|> tablesample (200b) repeatable (0) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0015", + "messageParameters" : { + "msg" : "byteLengthLiteral" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 44, + "fragment" : "tablesample (200b) repeatable (0)" + } ] +} + + +-- !query +table t +|> tablesample (200) repeatable (0) +-- !query analysis +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0016", + "messageParameters" : { + "bytesStr" : "200" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 43, + "fragment" : "tablesample (200) repeatable (0)" + } ] +} + + -- !query drop table t -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql index 3aa01d472e83f..31748fe1125ab 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/pipe-operators.sql @@ -326,6 +326,55 @@ table courseSales for `year` in (2012, 2013) ); +-- Sampling operators: positive tests. +-------------------------------------- + +-- We will use the REPEATABLE clause and/or adjust the sampling options to either remove no rows or +-- all rows to help keep the tests deterministic. +table t +|> tablesample (100 percent) repeatable (0); + +table t +|> tablesample (2 rows) repeatable (0); + +table t +|> tablesample (bucket 1 out of 1) repeatable (0); + +table t +|> tablesample (100 percent) repeatable (0) +|> tablesample (5 rows) repeatable (0) +|> tablesample (bucket 1 out of 1) repeatable (0); + +-- Sampling operators: negative tests. +-------------------------------------- + +-- The sampling method is required. +table t +|> tablesample (); + +-- Negative sampling options are not supported. +table t +|> tablesample (-100 percent); + +table t +|> tablesample (-5 rows); + +-- The sampling method may not refer to attribute names from the input relation. +table t +|> tablesample (x rows); + +-- The bucket number is invalid. +table t +|> tablesample (bucket 2 out of 1); + +-- Byte literals are not supported. +table t +|> tablesample (200b) repeatable (0); + +-- Invalid byte literal syntax. +table t +|> tablesample (200) repeatable (0); + -- Cleanup. ----------- drop table t; diff --git a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out index 2c6abe2a277ad..78b610b0d97c6 100644 --- a/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/pipe-operators.sql.out @@ -861,6 +861,204 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +table t +|> tablesample (100 percent) repeatable (0) +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> tablesample (2 rows) repeatable (0) +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> tablesample (bucket 1 out of 1) repeatable (0) +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> tablesample (100 percent) repeatable (0) +|> tablesample (5 rows) repeatable (0) +|> tablesample (bucket 1 out of 1) repeatable (0) +-- !query schema +struct +-- !query output +0 abc +1 def + + +-- !query +table t +|> tablesample () +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0014", + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 25, + "fragment" : "tablesample ()" + } ] +} + + +-- !query +table t +|> tablesample (-100 percent) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0064", + "messageParameters" : { + "msg" : "Sampling fraction (-1.0) must be on interval [0, 1]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 37, + "fragment" : "tablesample (-100 percent)" + } ] +} + + +-- !query +table t +|> tablesample (-5 rows) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_LIMIT_LIKE_EXPRESSION.IS_NEGATIVE", + "sqlState" : "42K0E", + "messageParameters" : { + "expr" : "\"-5\"", + "name" : "limit", + "v" : "-5" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 25, + "stopIndex" : 26, + "fragment" : "-5" + } ] +} + + +-- !query +table t +|> tablesample (x rows) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "INVALID_LIMIT_LIKE_EXPRESSION.IS_UNFOLDABLE", + "sqlState" : "42K0E", + "messageParameters" : { + "expr" : "\"x\"", + "name" : "limit" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 25, + "stopIndex" : 25, + "fragment" : "x" + } ] +} + + +-- !query +table t +|> tablesample (bucket 2 out of 1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0064", + "messageParameters" : { + "msg" : "Sampling fraction (2.0) must be on interval [0, 1]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 42, + "fragment" : "tablesample (bucket 2 out of 1)" + } ] +} + + +-- !query +table t +|> tablesample (200b) repeatable (0) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0015", + "messageParameters" : { + "msg" : "byteLengthLiteral" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 44, + "fragment" : "tablesample (200b) repeatable (0)" + } ] +} + + +-- !query +table t +|> tablesample (200) repeatable (0) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.parser.ParseException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_0016", + "messageParameters" : { + "bytesStr" : "200" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 12, + "stopIndex" : 43, + "fragment" : "tablesample (200) repeatable (0)" + } ] +} + + -- !query drop table t -- !query schema diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala index 1111a65c6a526..c76d44a1b82cf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala @@ -928,6 +928,15 @@ class SparkSqlParserSuite extends AnalysisTest with SharedSparkSession { | earningsYear FOR year IN (`2012`, `2013`, `2014`) |) |""".stripMargin) + // Sampling operations + def checkSample(query: String): Unit = { + val plan: LogicalPlan = parser.parsePlan(query) + assert(plan.collectFirst(_.isInstanceOf[Sample]).nonEmpty) + assert(plan.containsAnyPattern(UNRESOLVED_RELATION, LOCAL_RELATION)) + } + checkSample("TABLE t |> TABLESAMPLE (50 PERCENT)") + checkSample("TABLE t |> TABLESAMPLE (5 ROWS)") + checkSample("TABLE t |> TABLESAMPLE (BUCKET 4 OUT OF 10)") } } } From d97acc17dd0bce476a1f44e7cce14e8d13d95a51 Mon Sep 17 00:00:00 2001 From: Dongjoon Hyun Date: Wed, 2 Oct 2024 14:20:49 -0700 Subject: [PATCH 189/189] [SPARK-49853][SQL][TESTS] Increase test timeout of `PythonForeachWriterSuite` to `60s` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What changes were proposed in this pull request? This PR aims to increase test timeout of `PythonForeachWriterSuite` to `60s`. ### Why are the changes needed? To stablize `PythonForeachWriterSuite` in GitHub Action MacOS 15 Runner. For the failed cases, the data is still under generation. - https://github.com/apache/spark/actions/runs/11132652698/job/30936988757 ``` - UnsafeRowBuffer: handles more data than memory *** FAILED *** The code passed to eventually never returned normally. Attempted 237 times over 20.075615666999997 seconds. Last failure message: ArraySeq(1, ..., 1815) did not equal Range$Inclusive(1, ..., 2000) ``` GitHub Runners have different spec and macOS has very limited resources among them. - https://docs.github.com/en/actions/using-github-hosted-runners/using-github-hosted-runners/about-github-hosted-runners#standard-github-hosted-runners-for-public-repositories | Virtual Machine | Processor (CPU) | Memory (RAM) | Storage (SSD) | Workflow label | | -- | -- | -- | -- | -- | | Linux | 4 | 16 GB | 14 GB | ubuntu-latest,ubuntu-24.04,ubuntu-22.04,ubuntu-20.04 | | macOS | 3 (M1) | 7 GB | 14 GB | macos-latest,macos-14, macos-15 [Beta] | ### Does this PR introduce _any_ user-facing change? No, this is a test-only change. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48319 from dongjoon-hyun/SPARK-49853. Authored-by: Dongjoon Hyun Signed-off-by: Dongjoon Hyun --- .../spark/sql/execution/python/PythonForeachWriterSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala index 3a8ce569d1ba9..a2d3318361837 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala @@ -99,7 +99,7 @@ class PythonForeachWriterSuite extends SparkFunSuite with Eventually with Mockit } private val iterator = buffer.iterator private val outputBuffer = new ArrayBuffer[Int] - private val testTimeout = timeout(20.seconds) + private val testTimeout = timeout(60.seconds) private val intProj = UnsafeProjection.create(Array[DataType](IntegerType)) private val thread = new Thread() { override def run(): Unit = {