allisonwang-db commented on code in PR #44305: URL: https://github.com/apache/spark/pull/44305#discussion_r1426808206
########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonDataSource.scala: ########## @@ -20,58 +20,200 @@ package org.apache.spark.sql.execution.python import java.io.{DataInputStream, DataOutputStream} import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ import net.razorvine.pickle.Pickler -import org.apache.spark.api.python.{PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} -import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PythonDataSource} +import org.apache.spark.JobArtifactSet +import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType, PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.PythonUDF import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE} +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.{BinaryType, DataType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ +/** + * Data Source V2 wrapper for Python Data Source. + */ +class PythonTableProvider(shortName: String) extends TableProvider { + private var dataSourceInPython: PythonDataSourceCreationResult = _ + private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) + private lazy val source: UserDefinedPythonDataSource = + SparkSession.active.sessionState.dataSourceManager.lookupDataSource(shortName) + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + if (dataSourceInPython == null) { + dataSourceInPython = source.createDataSourceInPython(shortName, options, None) + } + dataSourceInPython.schema + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: java.util.Map[String, String]): Table = { + assert(partitioning.isEmpty) Review Comment: Do we need this assertion? Or shall we throw an exception when partitioning is not empty? What if a user creates a create table with partitioning: ``` CREATE TABLE test (i INT, j INT) USING my-python-data-source PARTITIONED BY (i) ``` ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonDataSource.scala: ########## @@ -20,58 +20,200 @@ package org.apache.spark.sql.execution.python import java.io.{DataInputStream, DataOutputStream} import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ import net.razorvine.pickle.Pickler -import org.apache.spark.api.python.{PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} -import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PythonDataSource} +import org.apache.spark.JobArtifactSet +import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType, PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.PythonUDF import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE} +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.{BinaryType, DataType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ +/** + * Data Source V2 wrapper for Python Data Source. + */ +class PythonTableProvider(shortName: String) extends TableProvider { + private var dataSourceInPython: PythonDataSourceCreationResult = _ + private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) + private lazy val source: UserDefinedPythonDataSource = + SparkSession.active.sessionState.dataSourceManager.lookupDataSource(shortName) + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + if (dataSourceInPython == null) { + dataSourceInPython = source.createDataSourceInPython(shortName, options, None) + } + dataSourceInPython.schema + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: java.util.Map[String, String]): Table = { + assert(partitioning.isEmpty) + val outputSchema = schema + new Table with SupportsRead { + override def name(): String = shortName + + override def capabilities(): java.util.Set[TableCapability] = java.util.EnumSet.of( + BATCH_READ, BATCH_WRITE) Review Comment: BATCH_WRITE is currently not supported ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/MapInBatchEvaluatorFactory.scala: ########## @@ -36,7 +36,7 @@ class MapInBatchEvaluatorFactory( sessionLocalTimeZone: String, largeVarTypes: Boolean, pythonRunnerConf: Map[String, String], - pythonMetrics: Map[String, SQLMetric], + pythonMetrics: Option[Map[String, SQLMetric]], Review Comment: Why do we need to change this to Optional? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonDataSource.scala: ########## @@ -20,58 +20,200 @@ package org.apache.spark.sql.execution.python import java.io.{DataInputStream, DataOutputStream} import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ import net.razorvine.pickle.Pickler -import org.apache.spark.api.python.{PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} -import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} -import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PythonDataSource} +import org.apache.spark.JobArtifactSet +import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType, PythonFunction, PythonWorkerUtils, SimplePythonFunction, SpecialLengths} +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.PythonUDF import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE} +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.{BinaryType, DataType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ +/** + * Data Source V2 wrapper for Python Data Source. + */ +class PythonTableProvider(shortName: String) extends TableProvider { + private var dataSourceInPython: PythonDataSourceCreationResult = _ + private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) + private lazy val source: UserDefinedPythonDataSource = + SparkSession.active.sessionState.dataSourceManager.lookupDataSource(shortName) + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + if (dataSourceInPython == null) { + dataSourceInPython = source.createDataSourceInPython(shortName, options, None) + } + dataSourceInPython.schema + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: java.util.Map[String, String]): Table = { + assert(partitioning.isEmpty) + val outputSchema = schema + new Table with SupportsRead { Review Comment: We can create a new class `PythonTable` to make it more extensible in the future. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org