[Spark] Delta Connect python client implementation + tests (ported from the branch-4.0-preview) (#4514)

<!--
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
-->

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

## Description

Based off of #4513

Adds the Delta Connect python client from `branch-4.0-preview1` to
master since we now are able to run our python tests using the latest
Spark 4.0 RC.

<!--
- 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.
-->

## How was this patch tested?

<!--
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?

<!--
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'.
-->

---------

Co-authored-by: Thang Long VU <[email protected]>

Author: allisonport-db

.github/workflows/spark_master_python_test.yaml

@@ -65,6 +65,7 @@ jobs:
           pipenv run pip install pip==24.0 setuptools==69.5.1 wheel==0.43.0
           pipenv run pip install flake8==3.9.0
           pipenv run pip install black==23.12.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
@@ -74,8 +75,15 @@ jobs:
           pipenv run pip install pydocstyle==3.0.0
           pipenv run pip install pandas==2.2.0
           pipenv run pip install pyarrow==11.0.0
-          pipenv run pip install numpy==1.21
-          pipenv run pip install https://dist.apache.org/repos/dist/dev/spark/v4.0.0-rc4-bin/pyspark-4.0.0.tar.gz
+          pipenv run pip install pypandoc==1.3.3
+          pipenv run pip install numpy==1.22.4
+          pipenv run pip install grpcio==1.67.0
+          pipenv run pip install grpcio-status==1.67.0
+          pipenv run pip install googleapis-common-protos==1.65.0
+          pipenv run pip install protobuf==5.29.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 https://dist.apache.org/repos/dist/dev/spark/v4.0.0-rc6-bin/pyspark-4.0.0.tar.gz
         if: steps.git-diff.outputs.diff
       - name: Run Python tests
         # when changing TEST_PARALLELISM_COUNT make sure to also change it in spark_master_test.yaml

dev/delta-connect-gen-protos.sh

@@ -71,8 +71,10 @@ for f in `find gen/proto/python/delta/connect -name "*.py*"`; do
   if [[ $f == *_pb2.py || $f == *_pb2_grpc.py ]]; then
     sed \
       -e 's/import spark.connect./import pyspark.sql.connect.proto./g' \
-      -e 's/from delta.connect import/from delta.connect.proto import/g' \
+      -e "s/DESCRIPTOR, 'spark.connect/DESCRIPTOR, 'pyspark.sql.connect.proto/g" \
+      -e 's/from spark.connect import/from pyspark.sql.connect.proto import/g' \
       -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

@@ -3,7 +3,12 @@ mypy==1.8.0
 flake8==3.9.0
 
 # Code Formatter
-black==23.9.1
+black==23.12.1
+
+# Spark Connect (required)
+grpcio>=1.67.0
+grpcio-status>=1.67.0
+googleapis-common-protos>=1.65.0
 
 # 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/relations_pb2.py

@@ -28,9 +28,9 @@
 
 
 from delta.connect.proto import base_pb2 as delta_dot_connect_dot_base__pb2
-from spark.connect import expressions_pb2 as spark_dot_connect_dot_expressions__pb2
-from spark.connect import relations_pb2 as spark_dot_connect_dot_relations__pb2
-from spark.connect import types_pb2 as spark_dot_connect_dot_types__pb2
+from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2
+from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2
+from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__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.
+#