You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2024/07/04 08:39:39 UTC
(spark-connect-go) branch master updated: [SPARK-48754] Separate session code from sql module
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-go.git
The following commit(s) were added to refs/heads/master by this push:
new f22f86b [SPARK-48754] Separate session code from sql module
f22f86b is described below
commit f22f86bd3290d440bba2e660b543f23a1a448f24
Author: Mathias Schwarz <16...@users.noreply.github.com>
AuthorDate: Thu Jul 4 17:39:32 2024 +0900
[SPARK-48754] Separate session code from sql module
### What changes were proposed in this pull request?
This pulls the session code into a separate package, enforcing the layered structure that is defined in the README. This is the change that we prepared by introducing the sparkExecutor interface earlier
### Why are the changes needed?
This improves the package structure by making it reflect the dependencies described in the README directly
### Does this PR introduce _any_ user-facing change?
It moves the session code to a new package but introduces no behavioral changes
### How was this patch tested?
Ran existing unit tests
Closes #27 from mathiasschw-db/mathiasschw-db/separate_session_from_sql.
Authored-by: Mathias Schwarz <16...@users.noreply.github.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
client/mocks/mocks.go | 47 +++++
client/{sql => session}/sparksession.go | 23 +--
client/session/sparksession_test.go | 235 ++++++++++++++++++++++++
client/sql/dataframe.go | 15 +-
client/sql/dataframe_test.go | 63 -------
client/sql/dataframereader.go | 7 +-
client/sql/dataframereader_test.go | 2 +-
client/sql/dataframewriter.go | 12 +-
client/sql/dataframewriter_test.go | 11 +-
client/sql/executeplanclient.go | 8 +-
client/sql/mocks_test.go | 94 +---------
client/sql/sparksession_test.go | 111 -----------
cmd/spark-connect-example-spark-session/main.go | 3 +-
13 files changed, 328 insertions(+), 303 deletions(-)
diff --git a/client/mocks/mocks.go b/client/mocks/mocks.go
new file mode 100644
index 0000000..3281243
--- /dev/null
+++ b/client/mocks/mocks.go
@@ -0,0 +1,47 @@
+package mocks
+
+import (
+ "context"
+
+ proto "github.com/apache/spark-connect-go/v3.5/internal/generated"
+ "google.golang.org/grpc/metadata"
+)
+
+type ProtoClient struct {
+ RecvResponse *proto.ExecutePlanResponse
+ RecvResponses []*proto.ExecutePlanResponse
+
+ Err error
+}
+
+func (p *ProtoClient) Recv() (*proto.ExecutePlanResponse, error) {
+ if len(p.RecvResponses) != 0 {
+ p.RecvResponse = p.RecvResponses[0]
+ p.RecvResponses = p.RecvResponses[1:]
+ }
+ return p.RecvResponse, p.Err
+}
+
+func (p *ProtoClient) Header() (metadata.MD, error) {
+ return nil, p.Err
+}
+
+func (p *ProtoClient) Trailer() metadata.MD {
+ return nil
+}
+
+func (p *ProtoClient) CloseSend() error {
+ return p.Err
+}
+
+func (p *ProtoClient) Context() context.Context {
+ return nil
+}
+
+func (p *ProtoClient) SendMsg(m interface{}) error {
+ return p.Err
+}
+
+func (p *ProtoClient) RecvMsg(m interface{}) error {
+ return p.Err
+}
diff --git a/client/sql/sparksession.go b/client/session/sparksession.go
similarity index 87%
rename from client/sql/sparksession.go
rename to client/session/sparksession.go
index 8121e9a..6176e06 100644
--- a/client/sql/sparksession.go
+++ b/client/session/sparksession.go
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-package sql
+package session
import (
"context"
@@ -22,14 +22,15 @@ import (
"github.com/apache/spark-connect-go/v35/client/channel"
"github.com/apache/spark-connect-go/v35/client/sparkerrors"
+ "github.com/apache/spark-connect-go/v35/client/sql"
proto "github.com/apache/spark-connect-go/v35/internal/generated"
"github.com/google/uuid"
"google.golang.org/grpc/metadata"
)
type SparkSession interface {
- Read() DataFrameReader
- Sql(ctx context.Context, query string) (DataFrame, error)
+ Read() sql.DataFrameReader
+ Sql(ctx context.Context, query string) (sql.DataFrame, error)
Stop() error
}
@@ -80,11 +81,11 @@ type sparkSessionImpl struct {
metadata metadata.MD
}
-func (s *sparkSessionImpl) Read() DataFrameReader {
- return newDataframeReader(s)
+func (s *sparkSessionImpl) Read() sql.DataFrameReader {
+ return sql.NewDataframeReader(s)
}
-func (s *sparkSessionImpl) Sql(ctx context.Context, query string) (DataFrame, error) {
+func (s *sparkSessionImpl) Sql(ctx context.Context, query string) (sql.DataFrame, error) {
plan := &proto.Plan{
OpType: &proto.Plan_Command{
Command: &proto.Command{
@@ -96,7 +97,7 @@ func (s *sparkSessionImpl) Sql(ctx context.Context, query string) (DataFrame, er
},
},
}
- responseClient, err := s.executePlan(ctx, plan)
+ responseClient, err := s.ExecutePlan(ctx, plan)
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to execute sql: %s: %w", query, err), sparkerrors.ExecutionError)
}
@@ -109,7 +110,7 @@ func (s *sparkSessionImpl) Sql(ctx context.Context, query string) (DataFrame, er
if sqlCommandResult == nil {
continue
}
- return newDataFrame(s, sqlCommandResult.GetRelation()), nil
+ return sql.NewDataFrame(s, sqlCommandResult.GetRelation()), nil
}
}
@@ -117,7 +118,7 @@ func (s *sparkSessionImpl) Stop() error {
return nil
}
-func (s *sparkSessionImpl) executePlan(ctx context.Context, plan *proto.Plan) (*executePlanClient, error) {
+func (s *sparkSessionImpl) ExecutePlan(ctx context.Context, plan *proto.Plan) (*sql.ExecutePlanClient, error) {
request := proto.ExecutePlanRequest{
SessionId: s.sessionId,
Plan: plan,
@@ -131,10 +132,10 @@ func (s *sparkSessionImpl) executePlan(ctx context.Context, plan *proto.Plan) (*
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to call ExecutePlan in session %s: %w", s.sessionId, err), sparkerrors.ExecutionError)
}
- return newExecutePlanClient(client), nil
+ return sql.NewExecutePlanClient(client), nil
}
-func (s *sparkSessionImpl) analyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
+func (s *sparkSessionImpl) AnalyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
request := proto.AnalyzePlanRequest{
SessionId: s.sessionId,
Analyze: &proto.AnalyzePlanRequest_Schema_{
diff --git a/client/session/sparksession_test.go b/client/session/sparksession_test.go
new file mode 100644
index 0000000..1e80f31
--- /dev/null
+++ b/client/session/sparksession_test.go
@@ -0,0 +1,235 @@
+package session
+
+import (
+ "bytes"
+ "context"
+ "testing"
+
+ "github.com/apache/arrow/go/v12/arrow"
+ "github.com/apache/arrow/go/v12/arrow/array"
+ "github.com/apache/arrow/go/v12/arrow/ipc"
+ "github.com/apache/arrow/go/v12/arrow/memory"
+ "github.com/apache/spark-connect-go/v3.5/client/mocks"
+ "github.com/apache/spark-connect-go/v3.5/client/sparkerrors"
+ "github.com/apache/spark-connect-go/v3.5/client/sql"
+ proto "github.com/apache/spark-connect-go/v3.5/internal/generated"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "google.golang.org/grpc"
+)
+
+type connectServiceClient struct {
+ t *testing.T
+
+ analysePlanResponse *proto.AnalyzePlanResponse
+ executePlanClient proto.SparkConnectService_ExecutePlanClient
+ expectedExecutePlanRequest *proto.ExecutePlanRequest
+
+ err error
+}
+
+func (c *connectServiceClient) ExecutePlan(ctx context.Context, in *proto.ExecutePlanRequest, opts ...grpc.CallOption) (proto.SparkConnectService_ExecutePlanClient, error) {
+ if c.expectedExecutePlanRequest != nil {
+ assert.Equal(c.t, c.expectedExecutePlanRequest, in)
+ }
+ return c.executePlanClient, c.err
+}
+
+func (c *connectServiceClient) AnalyzePlan(ctx context.Context, in *proto.AnalyzePlanRequest, opts ...grpc.CallOption) (*proto.AnalyzePlanResponse, error) {
+ return c.analysePlanResponse, c.err
+}
+
+func (c *connectServiceClient) Config(ctx context.Context, in *proto.ConfigRequest, opts ...grpc.CallOption) (*proto.ConfigResponse, error) {
+ return nil, c.err
+}
+
+func (c *connectServiceClient) AddArtifacts(ctx context.Context, opts ...grpc.CallOption) (proto.SparkConnectService_AddArtifactsClient, error) {
+ return nil, c.err
+}
+
+func (c *connectServiceClient) ArtifactStatus(ctx context.Context, in *proto.ArtifactStatusesRequest, opts ...grpc.CallOption) (*proto.ArtifactStatusesResponse, error) {
+ return nil, c.err
+}
+
+func (c *connectServiceClient) Interrupt(ctx context.Context, in *proto.InterruptRequest, opts ...grpc.CallOption) (*proto.InterruptResponse, error) {
+ return nil, c.err
+}
+
+func (c *connectServiceClient) ReattachExecute(ctx context.Context, in *proto.ReattachExecuteRequest, opts ...grpc.CallOption) (proto.SparkConnectService_ReattachExecuteClient, error) {
+ return nil, c.err
+}
+
+func (c *connectServiceClient) ReleaseExecute(ctx context.Context, in *proto.ReleaseExecuteRequest, opts ...grpc.CallOption) (*proto.ReleaseExecuteResponse, error) {
+ return nil, c.err
+}
+
+func TestAnalyzePlanCallsAnalyzePlanOnClient(t *testing.T) {
+ ctx := context.Background()
+ reponse := &proto.AnalyzePlanResponse{}
+ session := &sparkSessionImpl{
+ client: &connectServiceClient{
+ analysePlanResponse: reponse,
+ },
+ }
+ resp, err := session.AnalyzePlan(ctx, &proto.Plan{})
+ assert.NoError(t, err)
+ assert.NotNil(t, resp)
+}
+
+func TestAnalyzePlanFailsIfClientFails(t *testing.T) {
+ ctx := context.Background()
+ session := &sparkSessionImpl{
+ client: &connectServiceClient{
+ err: assert.AnError,
+ },
+ }
+ resp, err := session.AnalyzePlan(ctx, &proto.Plan{})
+ assert.Nil(t, resp)
+ assert.Error(t, err)
+}
+
+func TestExecutePlanCallsExecutePlanOnClient(t *testing.T) {
+ ctx := context.Background()
+
+ plan := &proto.Plan{}
+ request := &proto.ExecutePlanRequest{
+ Plan: plan,
+ UserContext: &proto.UserContext{
+ UserId: "na",
+ },
+ }
+ session := &sparkSessionImpl{
+ client: &connectServiceClient{
+ executePlanClient: &sql.ExecutePlanClient{},
+ expectedExecutePlanRequest: request,
+ t: t,
+ },
+ }
+ resp, err := session.ExecutePlan(ctx, plan)
+ assert.NoError(t, err)
+ assert.NotNil(t, resp)
+}
+
+func TestSQLCallsExecutePlanWithSQLOnClient(t *testing.T) {
+ ctx := context.Background()
+
+ query := "select * from bla"
+ plan := &proto.Plan{
+ OpType: &proto.Plan_Command{
+ Command: &proto.Command{
+ CommandType: &proto.Command_SqlCommand{
+ SqlCommand: &proto.SqlCommand{
+ Sql: query,
+ },
+ },
+ },
+ },
+ }
+ request := &proto.ExecutePlanRequest{
+ Plan: plan,
+ UserContext: &proto.UserContext{
+ UserId: "na",
+ },
+ }
+ session := &sparkSessionImpl{
+ client: &connectServiceClient{
+ executePlanClient: &sql.ExecutePlanClient{&mocks.ProtoClient{
+ RecvResponse: &proto.ExecutePlanResponse{
+ ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
+ SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
+ },
+ },
+ }},
+ expectedExecutePlanRequest: request,
+ t: t,
+ },
+ }
+ resp, err := session.Sql(ctx, query)
+ assert.NoError(t, err)
+ assert.NotNil(t, resp)
+}
+
+func TestNewSessionBuilderCreatesASession(t *testing.T) {
+ ctx := context.Background()
+ spark, err := NewSessionBuilder().Remote("sc:connection").Build(ctx)
+ assert.NoError(t, err)
+ assert.NotNil(t, spark)
+}
+
+func TestNewSessionBuilderFailsIfConnectionStringIsInvalid(t *testing.T) {
+ ctx := context.Background()
+ spark, err := NewSessionBuilder().Remote("invalid").Build(ctx)
+ assert.Error(t, err)
+ assert.ErrorIs(t, err, sparkerrors.InvalidInputError)
+ assert.Nil(t, spark)
+}
+
+func TestWriteResultStreamsArrowResultToCollector(t *testing.T) {
+ ctx := context.Background()
+
+ arrowFields := []arrow.Field{
+ {
+ Name: "show_string",
+ Type: &arrow.StringType{},
+ },
+ }
+ arrowSchema := arrow.NewSchema(arrowFields, nil)
+ var buf bytes.Buffer
+ arrowWriter := ipc.NewWriter(&buf, ipc.WithSchema(arrowSchema))
+ defer arrowWriter.Close()
+
+ alloc := memory.NewGoAllocator()
+ recordBuilder := array.NewRecordBuilder(alloc, arrowSchema)
+ defer recordBuilder.Release()
+
+ recordBuilder.Field(0).(*array.StringBuilder).Append("str1a\nstr1b")
+ recordBuilder.Field(0).(*array.StringBuilder).Append("str2")
+
+ record := recordBuilder.NewRecord()
+ defer record.Release()
+
+ err := arrowWriter.Write(record)
+ require.Nil(t, err)
+
+ query := "select * from bla"
+
+ session := &sparkSessionImpl{
+ client: &connectServiceClient{
+ executePlanClient: &sql.ExecutePlanClient{&mocks.ProtoClient{
+ RecvResponses: []*proto.ExecutePlanResponse{
+ {
+ ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
+ SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
+ },
+ },
+ {
+ ResponseType: &proto.ExecutePlanResponse_ArrowBatch_{
+ ArrowBatch: &proto.ExecutePlanResponse_ArrowBatch{
+ RowCount: 1,
+ Data: buf.Bytes(),
+ },
+ },
+ },
+ }},
+ },
+ t: t,
+ },
+ }
+ resp, err := session.Sql(ctx, query)
+ assert.NoError(t, err)
+ assert.NotNil(t, resp)
+ writer, err := resp.Repartition(1, []string{"1"})
+ assert.NoError(t, err)
+ collector := &testCollector{}
+ err = writer.WriteResult(ctx, collector, 1, false)
+ assert.NoError(t, err)
+ assert.Equal(t, []any{"str2"}, collector.row)
+}
+
+type testCollector struct {
+ row []any
+}
+
+func (t *testCollector) WriteRow(values []any) {
+ t.row = values
+}
diff --git a/client/sql/dataframe.go b/client/sql/dataframe.go
index 6870b49..8fa3897 100644
--- a/client/sql/dataframe.go
+++ b/client/sql/dataframe.go
@@ -66,11 +66,12 @@ type RangePartitionColumn struct {
// dataFrameImpl is an implementation of DataFrame interface.
type dataFrameImpl struct {
- sparkExecutor sparkExecutor
+ sparkExecutor SparkExecutor
relation *proto.Relation // TODO change to proto.Plan?
}
-func newDataFrame(sparkExecutor sparkExecutor, relation *proto.Relation) DataFrame {
+// NewDataFrame creates a new DataFrame
+func NewDataFrame(sparkExecutor SparkExecutor, relation *proto.Relation) DataFrame {
return &dataFrameImpl{
sparkExecutor: sparkExecutor,
relation: relation,
@@ -113,7 +114,7 @@ func (df *dataFrameImpl) WriteResult(ctx context.Context, collector ResultCollec
},
}
- responseClient, err := df.sparkExecutor.executePlan(ctx, plan)
+ responseClient, err := df.sparkExecutor.ExecutePlan(ctx, plan)
if err != nil {
return sparkerrors.WithType(fmt.Errorf("failed to show dataframe: %w", err), sparkerrors.ExecutionError)
}
@@ -136,7 +137,7 @@ func (df *dataFrameImpl) WriteResult(ctx context.Context, collector ResultCollec
}
func (df *dataFrameImpl) Schema(ctx context.Context) (*StructType, error) {
- response, err := df.sparkExecutor.analyzePlan(ctx, df.createPlan())
+ response, err := df.sparkExecutor.AnalyzePlan(ctx, df.createPlan())
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to analyze plan: %w", err), sparkerrors.ExecutionError)
}
@@ -146,7 +147,7 @@ func (df *dataFrameImpl) Schema(ctx context.Context) (*StructType, error) {
}
func (df *dataFrameImpl) Collect(ctx context.Context) ([]Row, error) {
- responseClient, err := df.sparkExecutor.executePlan(ctx, df.createPlan())
+ responseClient, err := df.sparkExecutor.ExecutePlan(ctx, df.createPlan())
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to execute plan: %w", err), sparkerrors.ExecutionError)
}
@@ -214,7 +215,7 @@ func (df *dataFrameImpl) CreateTempView(ctx context.Context, viewName string, re
},
}
- responseClient, err := df.sparkExecutor.executePlan(ctx, plan)
+ responseClient, err := df.sparkExecutor.ExecutePlan(ctx, plan)
if err != nil {
return sparkerrors.WithType(fmt.Errorf("failed to create temp view %s: %w", viewName, err), sparkerrors.ExecutionError)
}
@@ -302,7 +303,7 @@ func (df *dataFrameImpl) repartitionByExpressions(numPartitions int, partitionEx
},
},
}
- return newDataFrame(df.sparkExecutor, newRelation), nil
+ return NewDataFrame(df.sparkExecutor, newRelation), nil
}
func showArrowBatch(arrowBatch *proto.ExecutePlanResponse_ArrowBatch, collector ResultCollector) error {
diff --git a/client/sql/dataframe_test.go b/client/sql/dataframe_test.go
index 42500a1..cccd703 100644
--- a/client/sql/dataframe_test.go
+++ b/client/sql/dataframe_test.go
@@ -18,7 +18,6 @@ package sql
import (
"bytes"
- "context"
"testing"
"github.com/apache/arrow/go/v12/arrow"
@@ -316,65 +315,3 @@ type testCollector struct {
func (t *testCollector) WriteRow(values []any) {
t.row = values
}
-
-func TestWriteResultStreamsArrowResultToCollector(t *testing.T) {
- ctx := context.Background()
-
- arrowFields := []arrow.Field{
- {
- Name: "show_string",
- Type: &arrow.StringType{},
- },
- }
- arrowSchema := arrow.NewSchema(arrowFields, nil)
- var buf bytes.Buffer
- arrowWriter := ipc.NewWriter(&buf, ipc.WithSchema(arrowSchema))
- defer arrowWriter.Close()
-
- alloc := memory.NewGoAllocator()
- recordBuilder := array.NewRecordBuilder(alloc, arrowSchema)
- defer recordBuilder.Release()
-
- recordBuilder.Field(0).(*array.StringBuilder).Append("str1a\nstr1b")
- recordBuilder.Field(0).(*array.StringBuilder).Append("str2")
-
- record := recordBuilder.NewRecord()
- defer record.Release()
-
- err := arrowWriter.Write(record)
- require.Nil(t, err)
-
- query := "select * from bla"
-
- session := &sparkSessionImpl{
- client: &connectServiceClient{
- executePlanClient: &executePlanClient{&protoClient{
- recvResponses: []*proto.ExecutePlanResponse{
- {
- ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
- SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
- },
- },
- {
- ResponseType: &proto.ExecutePlanResponse_ArrowBatch_{
- ArrowBatch: &proto.ExecutePlanResponse_ArrowBatch{
- RowCount: 1,
- Data: buf.Bytes(),
- },
- },
- },
- }},
- },
- t: t,
- },
- }
- resp, err := session.Sql(ctx, query)
- assert.NoError(t, err)
- assert.NotNil(t, resp)
- writer, err := resp.Repartition(1, []string{"1"})
- assert.NoError(t, err)
- collector := &testCollector{}
- err = writer.WriteResult(ctx, collector, 1, false)
- assert.NoError(t, err)
- assert.Equal(t, []any{"str2"}, collector.row)
-}
diff --git a/client/sql/dataframereader.go b/client/sql/dataframereader.go
index 01f7748..17cce98 100644
--- a/client/sql/dataframereader.go
+++ b/client/sql/dataframereader.go
@@ -14,11 +14,12 @@ type DataFrameReader interface {
// dataFrameReaderImpl is an implementation of DataFrameReader interface.
type dataFrameReaderImpl struct {
- sparkSession sparkExecutor
+ sparkSession SparkExecutor
formatSource string
}
-func newDataframeReader(session sparkExecutor) DataFrameReader {
+// NewDataframeReader creates a new DataFrameReader
+func NewDataframeReader(session SparkExecutor) DataFrameReader {
return &dataFrameReaderImpl{
sparkSession: session,
}
@@ -34,7 +35,7 @@ func (w *dataFrameReaderImpl) Load(path string) (DataFrame, error) {
if w.formatSource != "" {
format = w.formatSource
}
- return newDataFrame(w.sparkSession, toRelation(path, format)), nil
+ return NewDataFrame(w.sparkSession, toRelation(path, format)), nil
}
func toRelation(path string, format string) *proto.Relation {
diff --git a/client/sql/dataframereader_test.go b/client/sql/dataframereader_test.go
index c5e73c0..f52ac55 100644
--- a/client/sql/dataframereader_test.go
+++ b/client/sql/dataframereader_test.go
@@ -7,7 +7,7 @@ import (
)
func TestLoadCreatesADataFrame(t *testing.T) {
- reader := newDataframeReader(nil)
+ reader := NewDataframeReader(nil)
source := "source"
path := "path"
reader.Format(source)
diff --git a/client/sql/dataframewriter.go b/client/sql/dataframewriter.go
index ae714c4..d8d3110 100644
--- a/client/sql/dataframewriter.go
+++ b/client/sql/dataframewriter.go
@@ -19,12 +19,12 @@ type DataFrameWriter interface {
Save(ctx context.Context, path string) error
}
-type sparkExecutor interface {
- executePlan(ctx context.Context, plan *proto.Plan) (*executePlanClient, error)
- analyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error)
+type SparkExecutor interface {
+ ExecutePlan(ctx context.Context, plan *proto.Plan) (*ExecutePlanClient, error)
+ AnalyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error)
}
-func newDataFrameWriter(sparkExecutor sparkExecutor, relation *proto.Relation) DataFrameWriter {
+func newDataFrameWriter(sparkExecutor SparkExecutor, relation *proto.Relation) DataFrameWriter {
return &dataFrameWriterImpl{
sparkExecutor: sparkExecutor,
relation: relation,
@@ -33,7 +33,7 @@ func newDataFrameWriter(sparkExecutor sparkExecutor, relation *proto.Relation) D
// dataFrameWriterImpl is an implementation of DataFrameWriter interface.
type dataFrameWriterImpl struct {
- sparkExecutor sparkExecutor
+ sparkExecutor SparkExecutor
relation *proto.Relation
saveMode string
formatSource string
@@ -74,7 +74,7 @@ func (w *dataFrameWriterImpl) Save(ctx context.Context, path string) error {
},
},
}
- responseClient, err := w.sparkExecutor.executePlan(ctx, plan)
+ responseClient, err := w.sparkExecutor.ExecutePlan(ctx, plan)
if err != nil {
return err
}
diff --git a/client/sql/dataframewriter_test.go b/client/sql/dataframewriter_test.go
index 659a061..613f509 100644
--- a/client/sql/dataframewriter_test.go
+++ b/client/sql/dataframewriter_test.go
@@ -5,6 +5,7 @@ import (
"io"
"testing"
+ "github.com/apache/spark-connect-go/v35/client/mocks"
proto "github.com/apache/spark-connect-go/v35/internal/generated"
"github.com/stretchr/testify/assert"
)
@@ -38,8 +39,8 @@ func TestGetSaveMode(t *testing.T) {
func TestSaveExecutesWriteOperationUntilEOF(t *testing.T) {
relation := &proto.Relation{}
executor := &testExecutor{
- client: newExecutePlanClient(&protoClient{
- err: io.EOF,
+ client: NewExecutePlanClient(&mocks.ProtoClient{
+ Err: io.EOF,
}),
}
ctx := context.Background()
@@ -55,8 +56,8 @@ func TestSaveExecutesWriteOperationUntilEOF(t *testing.T) {
func TestSaveFailsIfAnotherErrorHappensWhenReadingStream(t *testing.T) {
relation := &proto.Relation{}
executor := &testExecutor{
- client: newExecutePlanClient(&protoClient{
- err: assert.AnError,
+ client: NewExecutePlanClient(&mocks.ProtoClient{
+ Err: assert.AnError,
}),
}
ctx := context.Background()
@@ -72,7 +73,7 @@ func TestSaveFailsIfAnotherErrorHappensWhenReadingStream(t *testing.T) {
func TestSaveFailsIfAnotherErrorHappensWhenExecuting(t *testing.T) {
relation := &proto.Relation{}
executor := &testExecutor{
- client: newExecutePlanClient(&protoClient{}),
+ client: NewExecutePlanClient(&mocks.ProtoClient{}),
err: assert.AnError,
}
ctx := context.Background()
diff --git a/client/sql/executeplanclient.go b/client/sql/executeplanclient.go
index d180e5b..ebc7dda 100644
--- a/client/sql/executeplanclient.go
+++ b/client/sql/executeplanclient.go
@@ -9,12 +9,12 @@ import (
proto "github.com/apache/spark-connect-go/v35/internal/generated"
)
-type executePlanClient struct {
+type ExecutePlanClient struct {
proto.SparkConnectService_ExecutePlanClient
}
-func newExecutePlanClient(responseClient proto.SparkConnectService_ExecutePlanClient) *executePlanClient {
- return &executePlanClient{
+func NewExecutePlanClient(responseClient proto.SparkConnectService_ExecutePlanClient) *ExecutePlanClient {
+ return &ExecutePlanClient{
responseClient,
}
}
@@ -23,7 +23,7 @@ func newExecutePlanClient(responseClient proto.SparkConnectService_ExecutePlanCl
// discard the returned data if there is no error. This is necessary for handling GRPC response for
// saving data frame, since such consuming will trigger Spark Connect Driver really saving data frame.
// If we do not consume the returned GRPC stream, Spark Connect Driver will not really save data frame.
-func (c *executePlanClient) consumeAll() error {
+func (c *ExecutePlanClient) consumeAll() error {
for {
_, err := c.Recv()
if err != nil {
diff --git a/client/sql/mocks_test.go b/client/sql/mocks_test.go
index f0bdf3e..ded00e0 100644
--- a/client/sql/mocks_test.go
+++ b/client/sql/mocks_test.go
@@ -2,111 +2,23 @@ package sql
import (
"context"
- "testing"
proto "github.com/apache/spark-connect-go/v35/internal/generated"
- "github.com/stretchr/testify/assert"
- "google.golang.org/grpc"
- "google.golang.org/grpc/metadata"
)
type testExecutor struct {
- client *executePlanClient
+ client *ExecutePlanClient
response *proto.AnalyzePlanResponse
err error
}
-func (t *testExecutor) executePlan(ctx context.Context, plan *proto.Plan) (*executePlanClient, error) {
+func (t *testExecutor) ExecutePlan(ctx context.Context, plan *proto.Plan) (*ExecutePlanClient, error) {
if t.err != nil {
return nil, t.err
}
return t.client, nil
}
-func (t *testExecutor) analyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
+func (t *testExecutor) AnalyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
return t.response, nil
}
-
-type protoClient struct {
- recvResponse *proto.ExecutePlanResponse
- recvResponses []*proto.ExecutePlanResponse
-
- err error
-}
-
-func (p *protoClient) Recv() (*proto.ExecutePlanResponse, error) {
- if len(p.recvResponses) != 0 {
- p.recvResponse = p.recvResponses[0]
- p.recvResponses = p.recvResponses[1:]
- }
- return p.recvResponse, p.err
-}
-
-func (p *protoClient) Header() (metadata.MD, error) {
- return nil, p.err
-}
-
-func (p *protoClient) Trailer() metadata.MD {
- return nil
-}
-
-func (p *protoClient) CloseSend() error {
- return p.err
-}
-
-func (p *protoClient) Context() context.Context {
- return nil
-}
-
-func (p *protoClient) SendMsg(m interface{}) error {
- return p.err
-}
-
-func (p *protoClient) RecvMsg(m interface{}) error {
- return p.err
-}
-
-type connectServiceClient struct {
- t *testing.T
-
- analysePlanResponse *proto.AnalyzePlanResponse
- executePlanClient proto.SparkConnectService_ExecutePlanClient
- expectedExecutePlanRequest *proto.ExecutePlanRequest
-
- err error
-}
-
-func (c *connectServiceClient) ExecutePlan(ctx context.Context, in *proto.ExecutePlanRequest, opts ...grpc.CallOption) (proto.SparkConnectService_ExecutePlanClient, error) {
- if c.expectedExecutePlanRequest != nil {
- assert.Equal(c.t, c.expectedExecutePlanRequest, in)
- }
- return c.executePlanClient, c.err
-}
-
-func (c *connectServiceClient) AnalyzePlan(ctx context.Context, in *proto.AnalyzePlanRequest, opts ...grpc.CallOption) (*proto.AnalyzePlanResponse, error) {
- return c.analysePlanResponse, c.err
-}
-
-func (c *connectServiceClient) Config(ctx context.Context, in *proto.ConfigRequest, opts ...grpc.CallOption) (*proto.ConfigResponse, error) {
- return nil, c.err
-}
-
-func (c *connectServiceClient) AddArtifacts(ctx context.Context, opts ...grpc.CallOption) (proto.SparkConnectService_AddArtifactsClient, error) {
- return nil, c.err
-}
-
-func (c *connectServiceClient) ArtifactStatus(ctx context.Context, in *proto.ArtifactStatusesRequest, opts ...grpc.CallOption) (*proto.ArtifactStatusesResponse, error) {
- return nil, c.err
-}
-
-func (c *connectServiceClient) Interrupt(ctx context.Context, in *proto.InterruptRequest, opts ...grpc.CallOption) (*proto.InterruptResponse, error) {
- return nil, c.err
-}
-
-func (c *connectServiceClient) ReattachExecute(ctx context.Context, in *proto.ReattachExecuteRequest, opts ...grpc.CallOption) (proto.SparkConnectService_ReattachExecuteClient, error) {
- return nil, c.err
-}
-
-func (c *connectServiceClient) ReleaseExecute(ctx context.Context, in *proto.ReleaseExecuteRequest, opts ...grpc.CallOption) (*proto.ReleaseExecuteResponse, error) {
- return nil, c.err
-}
diff --git a/client/sql/sparksession_test.go b/client/sql/sparksession_test.go
deleted file mode 100644
index fee707d..0000000
--- a/client/sql/sparksession_test.go
+++ /dev/null
@@ -1,111 +0,0 @@
-package sql
-
-import (
- "context"
- "testing"
-
- "github.com/apache/spark-connect-go/v35/client/sparkerrors"
- proto "github.com/apache/spark-connect-go/v35/internal/generated"
- "github.com/stretchr/testify/assert"
-)
-
-func TestAnalyzePlanCallsAnalyzePlanOnClient(t *testing.T) {
- ctx := context.Background()
- reponse := &proto.AnalyzePlanResponse{}
- session := &sparkSessionImpl{
- client: &connectServiceClient{
- analysePlanResponse: reponse,
- },
- }
- resp, err := session.analyzePlan(ctx, &proto.Plan{})
- assert.NoError(t, err)
- assert.NotNil(t, resp)
-}
-
-func TestAnalyzePlanFailsIfClientFails(t *testing.T) {
- ctx := context.Background()
- session := &sparkSessionImpl{
- client: &connectServiceClient{
- err: assert.AnError,
- },
- }
- resp, err := session.analyzePlan(ctx, &proto.Plan{})
- assert.Nil(t, resp)
- assert.Error(t, err)
-}
-
-func TestExecutePlanCallsExecutePlanOnClient(t *testing.T) {
- ctx := context.Background()
-
- plan := &proto.Plan{}
- request := &proto.ExecutePlanRequest{
- Plan: plan,
- UserContext: &proto.UserContext{
- UserId: "na",
- },
- }
- session := &sparkSessionImpl{
- client: &connectServiceClient{
- executePlanClient: &executePlanClient{},
- expectedExecutePlanRequest: request,
- t: t,
- },
- }
- resp, err := session.executePlan(ctx, plan)
- assert.NoError(t, err)
- assert.NotNil(t, resp)
-}
-
-func TestSQLCallsExecutePlanWithSQLOnClient(t *testing.T) {
- ctx := context.Background()
-
- query := "select * from bla"
- plan := &proto.Plan{
- OpType: &proto.Plan_Command{
- Command: &proto.Command{
- CommandType: &proto.Command_SqlCommand{
- SqlCommand: &proto.SqlCommand{
- Sql: query,
- },
- },
- },
- },
- }
- request := &proto.ExecutePlanRequest{
- Plan: plan,
- UserContext: &proto.UserContext{
- UserId: "na",
- },
- }
- session := &sparkSessionImpl{
- client: &connectServiceClient{
- executePlanClient: &executePlanClient{&protoClient{
- recvResponse: &proto.ExecutePlanResponse{
- ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
- SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
- },
- },
- }},
- expectedExecutePlanRequest: request,
- t: t,
- },
- }
- resp, err := session.Sql(ctx, query)
- assert.NoError(t, err)
- assert.NotNil(t, resp)
-}
-
-func TestNewSessionBuilderCreatesASession(t *testing.T) {
- ctx := context.Background()
- spark, err := NewSessionBuilder().Remote("sc:connection").Build(ctx)
- assert.NoError(t, err)
- assert.NotNil(t, spark)
-}
-
-func TestNewSessionBuilderFailsIfConnectionStringIsInvalid(t *testing.T) {
- ctx := context.Background()
- spark, err := NewSessionBuilder().Remote("invalid").Build(ctx)
- assert.Error(t, err)
- assert.ErrorIs(t, err, sparkerrors.InvalidInputError)
- assert.Nil(t, spark)
-}
diff --git a/cmd/spark-connect-example-spark-session/main.go b/cmd/spark-connect-example-spark-session/main.go
index 885a616..8a8a15d 100644
--- a/cmd/spark-connect-example-spark-session/main.go
+++ b/cmd/spark-connect-example-spark-session/main.go
@@ -21,6 +21,7 @@ import (
"flag"
"log"
+ "github.com/apache/spark-connect-go/v35/client/session"
"github.com/apache/spark-connect-go/v35/client/sql"
)
@@ -32,7 +33,7 @@ var (
func main() {
flag.Parse()
ctx := context.Background()
- spark, err := sql.NewSessionBuilder().Remote(*remote).Build(ctx)
+ spark, err := session.NewSessionBuilder().Remote(*remote).Build(ctx)
if err != nil {
log.Fatalf("Failed: %s", err)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org