[Spark] Add Delta Connect Python Client for First 4.0 Preview branch (#3201)

<!--
Thanks for sending a pull request!  Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://github.com/delta-io/delta/blob/master/CONTRIBUTING.md
2. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP]
Your PR title ...'.
  3. Be sure to keep the PR description updated to reflect all changes.
  4. Please write your PR title to summarize what this PR proposes.
5. If possible, provide a concise example to reproduce the issue for a
faster review.
6. If applicable, include the corresponding issue number in the PR title
and link it in the body.
-->

#### Which Delta project/connector is this regarding?
<!--
Please add the component selected below to the beginning of the pull
request title
For example: [Spark] Title of my pull request
-->

- [X] Spark
- [ ] Standalone
- [ ] Flink
- [ ] Kernel
- [ ] Other (fill in here)

## Description

<!--
- Describe what this PR changes.
- Describe why we need the change.
 
If this PR resolves an issue be sure to include "Resolves #XXX" to
correctly link and close the issue upon merge.
-->
[Delta Connect Python Client
PR](#3136) but merge into
branch-4.0-preview1 instead of master.

## How was this patch tested?
Added UTs.
<!--
If tests were added, say they were added here. Please make sure to test
the changes thoroughly including negative and positive cases if
possible.
If the changes were tested in any way other than unit tests, please
clarify how you tested step by step (ideally copy and paste-able, so
that other reviewers can test and check, and descendants can verify in
the future).
If the changes were not tested, please explain why.
-->

## Does this PR introduce _any_ user-facing changes?
No.
<!--
If yes, please clarify the previous behavior and the change this PR
proposes - provide the console output, description and/or an example to
show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change
compared to the released Delta Lake versions or within the unreleased
branches such as master.
If no, write 'No'.
-->

---------

Signed-off-by: Tai Le Manh <[email protected]>
Signed-off-by: Shawn Chang <[email protected]>
Co-authored-by: Scott Sandre <[email protected]>
Co-authored-by: Tai Le Manh <[email protected]>
Co-authored-by: Johan Lasperas <[email protected]>
Co-authored-by: Allison Portis <[email protected]>
Co-authored-by: Tom van Bussel <[email protected]>
Co-authored-by: Jiaheng Tang <[email protected]>
Co-authored-by: Sumeet Varma <[email protected]>
Co-authored-by: Venki Korukanti <[email protected]>
Co-authored-by: Avril Aysha <[email protected]>
Co-authored-by: Shawn Chang <[email protected]>
Co-authored-by: Shawn Chang <[email protected]>
Co-authored-by: Hao Jiang <[email protected]>

Author: 12 people

.github/workflows/spark_master_test.yaml

@@ -57,22 +57,31 @@ jobs:
           pipenv --python 3.9 install
           pipenv run pip install flake8==3.9.0
           pipenv run pip install black==23.9.1
+          pipenv run pip install importlib_metadata==3.10.0
           pipenv run pip install mypy==1.8.0
           pipenv run pip install mypy-protobuf==3.3.0
           pipenv run pip install cryptography==37.0.4
           pipenv run pip install twine==4.0.1
           pipenv run pip install wheel==0.33.4
           pipenv run pip install setuptools==41.1.0
           pipenv run pip install pydocstyle==3.0.0
-          pipenv run pip install pandas==1.4.4
-          pipenv run pip install pyarrow==8.0.0
+          pipenv run pip install pandas==2.0.0
+          pipenv run pip install pyarrow==10.0.0
+          pipenv run pip install pypandoc==1.3.3
           pipenv run pip install numpy==1.21
+          pipenv run pip install grpcio==1.62.0
+          pipenv run pip install grpcio-status==1.62.0
+          pipenv run pip install googleapis-common-protos==1.56.4
+          pipenv run pip install protobuf==4.25.1
+          pipenv run pip install googleapis-common-protos-stubs==2.2.0
+          pipenv run pip install grpc-stubs==1.24.11
           pipenv run pip install pyspark==4.0.0.dev1
       - name: Run Spark Master tests
         # when changing TEST_PARALLELISM_COUNT make sure to also change it in spark_test.yaml
         # NOTE: in this branch, the default sparkVersion is the SPARK_MASTER_VERSION
         run: |
           TEST_PARALLELISM_COUNT=2 build/sbt -DsparkVersion=master "++ ${{ matrix.scala }}" clean connectServer/test
           TEST_PARALLELISM_COUNT=2 build/sbt -DsparkVersion=master "++ ${{ matrix.scala }}" clean connectServer/assembly connectClient/test
+          TEST_PARALLELISM_COUNT=2 RUN_DELTA_CONNECT_TESTS=true pipenv run python run-tests.py
           TEST_PARALLELISM_COUNT=2 pipenv run python run-tests.py --group spark
         if: steps.git-diff.outputs.diff

dev/delta-connect-gen-protos.sh

@@ -71,6 +71,7 @@ for f in `find gen/proto/python/delta/connect -name "*.py*"`; do
   if [[ $f == *_pb2.py || $f == *_pb2_grpc.py ]]; then
     sed \
       -e "s/DESCRIPTOR, 'delta.connect/DESCRIPTOR, 'delta.connect.proto/g" \
+      -e 's/from delta.connect import/from delta.connect.proto import/g' \
       $f > $f.tmp
     mv $f.tmp $f
   elif [[ $f == *.pyi ]]; then

dev/requirements.txt

@@ -1,9 +1,14 @@
 # Linter
-mypy==0.982
+mypy==1.8.0
 flake8==3.9.0
 
 # Code Formatter
 black==23.9.1
 
+# Spark Connect (required)
+grpcio>=1.62.0
+grpcio-status>=1.62.0
+googleapis-common-protos>=1.56.4
+
 # Spark and Delta Connect python proto generation plugin (optional)
 mypy-protobuf==3.3.0

examples/python/delta_connect.py

@@ -0,0 +1,100 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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.
+#
+
+"""
+To run this example you must follow these steps:
+
+Requirements:
+- Using Java 17
+- Spark 4.0.0-preview1+
+- delta-spark (python package) 4.0.0rc1+ and pyspark 4.0.0.dev1+
+
+(1) Start a local Spark connect server using this command:
+sbin/start-connect-server.sh \
+  --packages org.apache.spark:spark-connect_2.13:4.0.0-preview1,io.delta:delta-connect-server_2.13:{DELTA_VERSION},io.delta:delta-spark_2.13:{DELTA_VERSION},com.google.protobuf:protobuf-java:3.25.1 \
+  --conf "spark.sql.extensions=io.delta.sql.DeltaSparkSessionExtension" \
+  --conf "spark.sql.catalog.spark_catalog=org.apache.spark.sql.delta.catalog.DeltaCatalog" \
+  --conf "spark.connect.extensions.relation.classes"="org.apache.spark.sql.connect.delta.DeltaRelationPlugin" \
+  --conf "spark.connect.extensions.command.classes"="org.apache.spark.sql.connect.delta.DeltaCommandPlugin"
+* Be sure to replace DELTA_VERSION with the version you are using
+
+(2) Set the SPARK_REMOTE environment variable to point to your local Spark server
+export SPARK_REMOTE="sc://localhost:15002"
+
+(3) Run this file i.e. python3 examples/python/delta_connect.py
+"""
+
+import os
+from pyspark.sql import SparkSession
+from delta.tables import DeltaTable
+import shutil
+
+filePath = "/tmp/delta_connect"
+tableName = "delta_connect_table"
+
+def assert_dataframe_equals(df1, df2):
+    assert(df1.collect().sort() == df2.collect().sort())
+
+def cleanup(spark):
+    shutil.rmtree(filePath, ignore_errors=True)
+    spark.sql(f"DROP TABLE IF EXISTS {tableName}")
+
+# --------------------- Set up Spark Connect spark session ------------------------
+
+assert os.getenv("SPARK_REMOTE"), "Must point to Spark Connect server using SPARK_REMOTE"
+
+spark = SparkSession.builder \
+    .appName("delta_connect") \
+    .remote(os.getenv("SPARK_REMOTE")) \
+    .getOrCreate()
+
+# Clean up any previous runs
+cleanup(spark)
+
+# -------------- Try reading non-existent table (should fail with an exception) ----------------
+
+# Using forPath
+try:
+    DeltaTable.forPath(spark, filePath).toDF().show()
+except Exception as e:
+    assert "DELTA_MISSING_DELTA_TABLE" in str(e)
+else:
+    assert False, "Expected exception to be thrown for missing table"
+
+# Using forName
+try:
+    DeltaTable.forName(spark, tableName).toDF().show()
+except Exception as e:
+    assert "DELTA_MISSING_DELTA_TABLE" in str(e)
+else:
+    assert False, "Expected exception to be thrown for missing table"
+
+# ------------------------ Write basic table and check that results match ----------------------
+
+# By table name
+spark.range(5).write.format("delta").saveAsTable(tableName)
+assert_dataframe_equals(DeltaTable.forName(spark, tableName).toDF(), spark.range(5))
+assert_dataframe_equals(spark.read.format("delta").table(tableName), spark.range(5))
+assert_dataframe_equals(spark.sql(f"SELECT * FROM {tableName}"), spark.range(5))
+
+# By table path
+spark.range(10).write.format("delta").save(filePath)
+assert_dataframe_equals(DeltaTable.forPath(spark, filePath).toDF(), spark.range(10))
+assert_dataframe_equals(spark.read.format("delta").load(filePath), spark.range(10))
+assert_dataframe_equals(spark.sql(f"SELECT * FROM delta.`{filePath}`"), spark.range(10))
+
+# ---------------------------------- Clean up ----------------------------------------
+cleanup(spark)

python/delta/connect/__init__.py

@@ -0,0 +1,19 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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 delta.connect.tables import DeltaTable
+
+__all__ = ['DeltaTable']

python/delta/connect/plan.py

@@ -0,0 +1,55 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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 Optional
+
+import delta.connect.proto as proto
+
+from pyspark.sql.connect.client import SparkConnectClient
+from pyspark.sql.connect.plan import LogicalPlan
+import pyspark.sql.connect.proto as spark_proto
+
+
+class DeltaLogicalPlan(LogicalPlan):
+    def __init__(self, child: Optional[LogicalPlan]) -> None:
+        super().__init__(child)
+
+    def plan(self, session: SparkConnectClient) -> spark_proto.Relation:
+        plan = spark_proto.Relation()
+        plan.extension.Pack(self.to_delta_relation(session))
+        return plan
+
+    def to_delta_relation(self, session: SparkConnectClient) -> proto.DeltaRelation:
+        ...
+
+    def command(self, session: SparkConnectClient) -> spark_proto.Command:
+        command = spark_proto.Command()
+        command.extension.Pack(self.to_delta_command(session))
+        return command
+
+    def to_delta_command(self, session: SparkConnectClient) -> proto.DeltaCommand:
+        ...
+
+
+class DeltaScan(DeltaLogicalPlan):
+    def __init__(self, table: proto.DeltaTable) -> None:
+        super().__init__(None)
+        self._table = table
+
+    def to_delta_relation(self, client: SparkConnectClient) -> proto.DeltaRelation:
+        relation = proto.DeltaRelation()
+        relation.scan.table.CopyFrom(self._table)
+        return relation

python/delta/connect/proto/__init__.py

@@ -0,0 +1,19 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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 delta.connect.proto.base_pb2 import *
+from delta.connect.proto.commands_pb2 import *
+from delta.connect.proto.relations_pb2 import *

python/delta/connect/proto/commands_pb2.py

@@ -27,7 +27,7 @@
 _sym_db = _symbol_database.Default()
 
 
-from delta.connect import base_pb2 as delta_dot_connect_dot_base__pb2
+from delta.connect.proto import base_pb2 as delta_dot_connect_dot_base__pb2
 
 
 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(

python/delta/connect/proto/relations_pb2.py

@@ -27,7 +27,7 @@
 _sym_db = _symbol_database.Default()
 
 
-from delta.connect import base_pb2 as delta_dot_connect_dot_base__pb2
+from delta.connect.proto import base_pb2 as delta_dot_connect_dot_base__pb2
 
 
 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(

python/delta/connect/tables.py

@@ -0,0 +1,88 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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, Optional
+
+from delta.connect.plan import DeltaScan
+import delta.connect.proto as proto
+from delta.tables import DeltaTable as LocalDeltaTable
+
+from pyspark.sql.connect.dataframe import DataFrame
+from pyspark.sql.connect.plan import LogicalPlan, SubqueryAlias
+from pyspark.sql.connect.session import SparkSession
+
+
+class DeltaTable(object):
+    def __init__(
+        self,
+        spark: SparkSession,
+        path: Optional[str] = None,
+        tableOrViewName: Optional[str] = None,
+        hadoopConf: Dict[str, str] = dict(),
+        plan: Optional[LogicalPlan] = None
+    ) -> None:
+        self._spark = spark
+        self._path = path
+        self._tableOrViewName = tableOrViewName
+        self._hadoopConf = hadoopConf
+        if plan is not None:
+            self._plan = plan
+        else:
+            self._plan = DeltaScan(self._to_proto())
+
+    def toDF(self) -> DataFrame:
+        return DataFrame(self._plan, session=self._spark)
+
+    def alias(self, aliasName: str) -> "DeltaTable":
+        return DeltaTable(
+            self._spark,
+            self._path,
+            self._tableOrViewName,
+            self._hadoopConf,
+            SubqueryAlias(self._plan, aliasName)
+        )
+
+    @classmethod
+    def forPath(
+        cls,
+        sparkSession: SparkSession,
+        path: str,
+        hadoopConf: Dict[str, str] = dict()
+    ) -> "DeltaTable":
+        assert sparkSession is not None
+        return DeltaTable(sparkSession, path=path, hadoopConf=hadoopConf)
+
+    @classmethod
+    def forName(
+        cls, sparkSession: SparkSession, tableOrViewName: str
+    ) -> "DeltaTable":
+        assert sparkSession is not None
+        return DeltaTable(sparkSession, tableOrViewName=tableOrViewName)
+
+    def _to_proto(self) -> proto.DeltaTable:
+        result = proto.DeltaTable()
+        if self._path is not None:
+            result.path.path = self._path
+        if self._tableOrViewName is not None:
+            result.table_or_view_name = self._tableOrViewName
+        return result
+
+
+DeltaTable.__doc__ = LocalDeltaTable.__doc__
+DeltaTable.toDF.__doc__ = LocalDeltaTable.toDF.__doc__
+DeltaTable.alias.__doc__ = LocalDeltaTable.alias.__doc__
+DeltaTable.forPath.__func__.__doc__ = LocalDeltaTable.forPath.__doc__
+DeltaTable.forName.__func__.__doc__ = LocalDeltaTable.forName.__doc__

python/delta/connect/testing/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright (2024) The Delta Lake Project Authors.
+#
+# Licensed 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.
+#