(spark-connect-go) branch master updated: Add Sample functionality in DataFrame

[mg...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mgrund 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 f9398aa  Add Sample functionality in DataFrame
f9398aa is described below

commit f9398aa124b5effaed70f9b6b804cb82b94d40fb
Author: vatsal <va...@gmail.com>
AuthorDate: Thu Jan 2 20:30:37 2025 +0100

    Add Sample functionality in DataFrame
    
    Signed-off-by: Vatsal <vatsal.v.anandgmail.com>
    
    ### What changes were proposed in this pull request?
    This PR adds support for the `sample` DF function.
    
    ### Why are the changes needed?
    Compatibility
    
    ### Does this PR introduce _any_ user-facing change?
    new DF functions
    
    ### How was this patch tested?
    Added e2e tests
    
    Closes #84 from imvtsl/issue.
    
    Lead-authored-by: vatsal <va...@gmail.com>
    Co-authored-by: Martin Grund <ma...@databricks.com>
    Signed-off-by: Martin Grund <ma...@databricks.com>
---
 internal/tests/integration/dataframe_test.go | 144 +++++++++++++++++++++++++++
 spark/sql/dataframe.go                       |  53 +++++++++-
 spark/sql/dataframestatfunctions.go          |  24 +++++
 3 files changed, 220 insertions(+), 1 deletion(-)

diff --git a/internal/tests/integration/dataframe_test.go b/internal/tests/integration/dataframe_test.go
index 6220b8e..720cc98 100644
--- a/internal/tests/integration/dataframe_test.go
+++ b/internal/tests/integration/dataframe_test.go
@@ -805,6 +805,150 @@ func TestDataFrame_WithOption(t *testing.T) {
 	assert.Equal(t, int64(10), c)
 }
 
+func TestDataFrame_Sample(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+	testCases := []struct {
+		name     string
+		fraction float64
+	}{
+		{
+			name:     "Default behavior",
+			fraction: 0.1,
+		},
+		{
+			name:     "Large fraction",
+			fraction: 0.9,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			sampledDF, err := df.Sample(ctx, tc.fraction)
+			assert.NoError(t, err)
+			count, err := sampledDF.Count(ctx)
+			assert.NoError(t, err)
+			expectedSize := int(100 * tc.fraction)
+			assert.InDelta(t, expectedSize, count, float64(expectedSize), 10)
+			rows, err := sampledDF.Collect(ctx)
+			assert.NoError(t, err)
+			// If sampling without replacement, check for duplicates
+			seen := make(map[int64]bool)
+			for _, row := range rows {
+				value := row.At(0).(int64)
+				if seen[value] {
+					t.Fatal("Found duplicate value when sampling without replacement")
+				}
+				seen[value] = true
+			}
+		})
+	}
+}
+
+func TestDataFrame_SampleWithReplacement(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+	testCases := []struct {
+		name            string
+		withReplacement bool
+		fraction        float64
+	}{
+		{
+			name:            "With replacement",
+			withReplacement: true,
+			fraction:        0.1,
+		},
+		{
+			name:            "Without replacement",
+			withReplacement: false,
+			fraction:        0.1,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			sampledDF, err := df.SampleWithReplacement(ctx, tc.withReplacement, tc.fraction)
+			assert.NoError(t, err)
+			count, err := sampledDF.Count(ctx)
+			assert.NoError(t, err)
+			expectedSize := int(100 * tc.fraction)
+			assert.InDelta(t, expectedSize, count, float64(expectedSize), 10)
+			rows, err := sampledDF.Collect(ctx)
+			assert.NoError(t, err)
+			// If sampling without replacement, check for duplicates
+			if tc.withReplacement == false {
+				seen := make(map[int64]bool)
+				for _, row := range rows {
+					value := row.At(0).(int64)
+					if seen[value] {
+						t.Fatal("Found duplicate value when sampling without replacement")
+					}
+					seen[value] = true
+				}
+			}
+		})
+	}
+}
+
+func TestDataFrame_SampleSeed(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+	fraction := 0.1
+	seed := int64(17)
+	sampledDF, err := df.SampleWithSeed(ctx, fraction, seed)
+	assert.NoError(t, err)
+	count, err := sampledDF.Count(ctx)
+	assert.NoError(t, err)
+	expectedSize := int(100 * fraction)
+	assert.InDelta(t, expectedSize, count, float64(expectedSize), 10)
+	rows, err := sampledDF.Collect(ctx)
+	assert.NoError(t, err)
+	// If sampling without replacement, check for duplicates
+	seen := make(map[int64]bool)
+	for _, row := range rows {
+		value := row.At(0).(int64)
+		if seen[value] {
+			t.Fatal("Found duplicate value when sampling without replacement")
+		}
+		seen[value] = true
+	}
+	// same seed should return same output
+	sampledDFRepeat, err := df.SampleWithSeed(ctx, fraction, seed)
+	assert.NoError(t, err)
+	count2, err := sampledDFRepeat.Count(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, count, count2)
+	rows2, err := sampledDFRepeat.Collect(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, rows, rows2)
+}
+
+func TestDataFrame_SampleWithReplacementSeed(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+	fraction := 0.1
+	seed := int64(17)
+	sampledDF, err := df.SampleWithReplacementAndSeed(ctx, true, fraction, seed)
+	assert.NoError(t, err)
+	count, err := sampledDF.Count(ctx)
+	assert.NoError(t, err)
+	expectedSize := int(100 * fraction)
+	assert.InDelta(t, expectedSize, count, float64(expectedSize), 10)
+	rows, err := sampledDF.Collect(ctx)
+	assert.NoError(t, err)
+	// same seed should return same output
+	sampledDFRepeat, err := df.SampleWithReplacementAndSeed(ctx, true, fraction, seed)
+	assert.NoError(t, err)
+	count2, err := sampledDFRepeat.Count(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, count, count2)
+	rows2, err := sampledDFRepeat.Collect(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, rows, rows2)
+}
+
 func TestDataFrame_Unpivot(t *testing.T) {
 	ctx, spark := connect()
 	data := [][]any{{1, 11, 1.1}, {2, 12, 1.2}}
diff --git a/spark/sql/dataframe.go b/spark/sql/dataframe.go
index 1d2b20b..ce67b0e 100644
--- a/spark/sql/dataframe.go
+++ b/spark/sql/dataframe.go
@@ -177,6 +177,14 @@ type DataFrame interface {
 	Rollup(ctx context.Context, cols ...column.Convertible) *GroupedData
 	// SameSemantics returns true if the other DataFrame has the same semantics.
 	SameSemantics(ctx context.Context, other DataFrame) (bool, error)
+	// Sample samples a data frame without replacement and random seed.
+	Sample(ctx context.Context, fraction float64) (DataFrame, error)
+	// SampleWithReplacement samples a data frame with random seed and with/without replacement.
+	SampleWithReplacement(ctx context.Context, withReplacement bool, fraction float64) (DataFrame, error)
+	// SampleWithSeed samples a data frame without replacement and given seed.
+	SampleWithSeed(ctx context.Context, fraction float64, seed int64) (DataFrame, error)
+	// SampleWithReplacementAndSeed samples a data frame with/without replacement and given seed.
+	SampleWithReplacementAndSeed(ctx context.Context, withReplacement bool, fraction float64, seed int64) (DataFrame, error)
 	// Show uses WriteResult to write the data frames to the console output.
 	Show(ctx context.Context, numRows int, truncate bool) error
 	// Schema returns the schema for the current data frame.
@@ -1401,6 +1409,50 @@ func (df *dataFrameImpl) Summary(ctx context.Context, statistics ...string) Data
 	return NewDataFrame(df.session, rel)
 }
 
+func (df *dataFrameImpl) Sample(ctx context.Context, fraction float64) (DataFrame, error) {
+	return df.sample(ctx, nil, fraction, nil)
+}
+
+func (df *dataFrameImpl) SampleWithReplacement(ctx context.Context, withReplacement bool, fraction float64) (DataFrame, error) {
+	return df.sample(ctx, &withReplacement, fraction, nil)
+}
+
+func (df *dataFrameImpl) SampleWithSeed(ctx context.Context, fraction float64, seed int64) (DataFrame, error) {
+	return df.sample(ctx, nil, fraction, &seed)
+}
+
+func (df *dataFrameImpl) SampleWithReplacementAndSeed(ctx context.Context, withReplacement bool, fraction float64, seed int64) (DataFrame, error) {
+	return df.sample(ctx, &withReplacement, fraction, &seed)
+}
+
+func (df *dataFrameImpl) sample(ctx context.Context, withReplacement *bool, fraction float64, seed *int64) (DataFrame, error) {
+	if seed == nil {
+		defaultSeed := rand.Int64()
+		seed = &defaultSeed
+	}
+
+	if withReplacement == nil {
+		defaultWithReplacement := false
+		withReplacement = &defaultWithReplacement
+	}
+
+	rel := &proto.Relation{
+		Common: &proto.RelationCommon{
+			PlanId: newPlanId(),
+		},
+		RelType: &proto.Relation_Sample{
+			Sample: &proto.Sample{
+				Input:           df.relation,
+				LowerBound:      0,
+				UpperBound:      fraction,
+				WithReplacement: withReplacement,
+				Seed:            seed,
+			},
+		},
+	}
+	return NewDataFrame(df.session, rel), nil
+}
+
 func (df *dataFrameImpl) Replace(ctx context.Context,
 	toReplace []types.PrimitiveTypeLiteral, values []types.PrimitiveTypeLiteral, cols ...string,
 ) (DataFrame, error) {
@@ -1441,7 +1493,6 @@ func (df *dataFrameImpl) Replace(ctx context.Context,
 		Common: &proto.RelationCommon{
 			PlanId: newPlanId(),
 		},
-
 		RelType: &proto.Relation_Replace{
 			Replace: &proto.NAReplace{
 				Input:        df.relation,
diff --git a/spark/sql/dataframestatfunctions.go b/spark/sql/dataframestatfunctions.go
index 7c05865..d68bfad 100644
--- a/spark/sql/dataframestatfunctions.go
+++ b/spark/sql/dataframestatfunctions.go
@@ -25,12 +25,36 @@ type DataFrameStatFunctions interface {
 	CrossTab(ctx context.Context, col1, col2 string) DataFrame
 	FreqItems(ctx context.Context, cols ...string) DataFrame
 	FreqItemsWithSupport(ctx context.Context, support float64, cols ...string) DataFrame
+	Sample(ctx context.Context, fraction float64) (DataFrame, error)
+	SampleWithReplacement(ctx context.Context, withReplacement bool, fraction float64) (DataFrame, error)
+	SampleWithSeed(ctx context.Context, fraction float64, seed int64) (DataFrame, error)
+	SampleWithReplacementAndSeed(ctx context.Context, withReplacement bool, fraction float64, seed int64) (DataFrame, error)
 }
 
 type dataFrameStatFunctionsImpl struct {
 	df DataFrame
 }
 
+func (d *dataFrameStatFunctionsImpl) Sample(ctx context.Context, fraction float64) (DataFrame, error) {
+	return d.df.Sample(ctx, fraction)
+}
+
+func (d *dataFrameStatFunctionsImpl) SampleWithReplacement(ctx context.Context,
+	withReplacement bool, fraction float64,
+) (DataFrame, error) {
+	return d.df.SampleWithReplacement(ctx, withReplacement, fraction)
+}
+
+func (d *dataFrameStatFunctionsImpl) SampleWithSeed(ctx context.Context, fraction float64, seed int64) (DataFrame, error) {
+	return d.df.SampleWithSeed(ctx, fraction, seed)
+}
+
+func (d *dataFrameStatFunctionsImpl) SampleWithReplacementAndSeed(ctx context.Context,
+	withReplacement bool, fraction float64, seed int64,
+) (DataFrame, error) {
+	return d.df.SampleWithReplacementAndSeed(ctx, withReplacement, fraction, seed)
+}
+
 func (d *dataFrameStatFunctionsImpl) ApproxQuantile(ctx context.Context, probabilities []float64,
 	relativeError float64, cols ...string,
 ) ([][]float64, error) {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org