fix: Use Go memory allocator for arrow (#810)

yevgenypats · web-flow · commit b54e5e16378d · 2023-04-20T05:33:54.000Z
Following this short discussion - apache/arrow#35232 It seems we don't really need to use Retain/Release outside the arrow library. We can always bring this back in the future if we would like to experiment if this brings better memory performance then the default go allocator.
diff --git a/internal/memdb/memdb.go b/internal/memdb/memdb.go
@@ -86,10 +86,8 @@ func (c *client) overwrite(table *arrow.Schema, data arrow.Record) {
 			}
 		}
 		if found {
-			tmp := c.memoryDB[tableName][i]
 			c.memoryDB[tableName] = append(c.memoryDB[tableName][:i], c.memoryDB[tableName][i+1:]...)
 			c.memoryDB[tableName] = append(c.memoryDB[tableName], data)
-			tmp.Release()
 			return
 		}
 	}
@@ -110,11 +108,6 @@ func (c *client) Migrate(_ context.Context, tables schema.Schemas) error {
 		if changes == nil {
 			continue
 		}
-		for _, t := range c.memoryDB {
-			for _, row := range t {
-				row.Release()
-			}
-		}
 		c.memoryDB[tableName] = make([]arrow.Record, 0)
 		c.tables[tableName] = table
 	}
@@ -204,11 +197,6 @@ func (*client) Metrics() destination.Metrics {
 }
 
 func (c *client) Close(context.Context) error {
-	for _, table := range c.memoryDB {
-		for _, row := range table {
-			row.Release()
-		}
-	}
 	c.memoryDB = nil
 	return nil
 }
@@ -230,8 +218,6 @@ func (c *client) deleteStaleTable(_ context.Context, table *arrow.Schema, source
 			rowSyncTime := row.Column(syncColIndex).(*array.Timestamp).Value(0).ToTime(arrow.Microsecond).UTC()
 			if !rowSyncTime.Before(syncTime) {
 				filteredTable = append(filteredTable, c.memoryDB[tableName][i])
-			} else {
-				c.memoryDB[tableName][i].Release()
 			}
 		}
 	}
diff --git a/internal/memdb/memdb_test.go b/internal/memdb/memdb_test.go
@@ -6,7 +6,6 @@ import (
 	"time"
 
 	"github.com/apache/arrow/go/v12/arrow"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v2/plugins/destination"
 	"github.com/cloudquery/plugin-sdk/v2/specs"
 	"github.com/cloudquery/plugin-sdk/v2/testdata"
@@ -122,17 +121,14 @@ func TestOnWriteError(t *testing.T) {
 	sourceSpec := specs.Source{
 		Name: sourceName,
 	}
-	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-	defer mem.AssertSize(t, 0)
 	ch := make(chan arrow.Record, 1)
 	opts := testdata.GenTestDataOptions{
 		SourceName: "test",
 		SyncTime:   time.Now(),
 		MaxRows:    1,
 		StableUUID: uuid.Nil,
 	}
-	record := testdata.GenTestData(mem, table.ToArrowSchema(), opts)[0]
-	defer record.Release()
+	record := testdata.GenTestData(table.ToArrowSchema(), opts)[0]
 	ch <- record
 	close(ch)
 	err := p.Write(ctx, sourceSpec, tables, syncTime, ch)
@@ -160,8 +156,6 @@ func TestOnWriteCtxCancelled(t *testing.T) {
 	sourceSpec := specs.Source{
 		Name: sourceName,
 	}
-	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-	defer mem.AssertSize(t, 0)
 	ch := make(chan arrow.Record, 1)
 	ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
 	opts := testdata.GenTestDataOptions{
@@ -170,8 +164,7 @@ func TestOnWriteCtxCancelled(t *testing.T) {
 		MaxRows:    1,
 		StableUUID: uuid.Nil,
 	}
-	record := testdata.GenTestData(mem, table.ToArrowSchema(), opts)[0]
-	defer record.Release()
+	record := testdata.GenTestData(table.ToArrowSchema(), opts)[0]
 	ch <- record
 	defer cancel()
 	err := p.Write(ctx, sourceSpec, tables, syncTime, ch)
diff --git a/plugins/destination/plugin_testing.go b/plugins/destination/plugin_testing.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/array"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v2/schema"
 	"github.com/cloudquery/plugin-sdk/v2/specs"
 	"github.com/cloudquery/plugin-sdk/v2/types"
@@ -116,11 +115,9 @@ func PluginTestSuiteRunner(t *testing.T, newPlugin NewPluginFunc, destSpec specs
 		if suite.tests.SkipOverwrite {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.Name = "test_write_overwrite"
 		p := newPlugin()
-		if err := suite.destinationPluginTestWriteOverwrite(ctx, mem, p, logger, destSpec); err != nil {
+		if err := suite.destinationPluginTestWriteOverwrite(ctx, p, logger, destSpec); err != nil {
 			t.Fatal(err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -133,11 +130,9 @@ func PluginTestSuiteRunner(t *testing.T, newPlugin NewPluginFunc, destSpec specs
 		if suite.tests.SkipOverwrite || suite.tests.SkipDeleteStale {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.Name = "test_write_overwrite_delete_stale"
 		p := newPlugin()
-		if err := suite.destinationPluginTestWriteOverwriteDeleteStale(ctx, mem, p, logger, destSpec); err != nil {
+		if err := suite.destinationPluginTestWriteOverwriteDeleteStale(ctx, p, logger, destSpec); err != nil {
 			t.Fatal(err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -150,36 +145,30 @@ func PluginTestSuiteRunner(t *testing.T, newPlugin NewPluginFunc, destSpec specs
 		if suite.tests.SkipMigrateOverwrite {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.WriteMode = specs.WriteModeOverwrite
 		destSpec.Name = "test_migrate_overwrite"
-		suite.destinationPluginTestMigrate(ctx, mem, t, newPlugin, logger, destSpec, tests.MigrateStrategyOverwrite)
+		suite.destinationPluginTestMigrate(ctx, t, newPlugin, logger, destSpec, tests.MigrateStrategyOverwrite)
 	})
 
 	t.Run("TestMigrateOverwriteForce", func(t *testing.T) {
 		t.Helper()
 		if suite.tests.SkipMigrateOverwriteForce {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.WriteMode = specs.WriteModeOverwrite
 		destSpec.MigrateMode = specs.MigrateModeForced
 		destSpec.Name = "test_migrate_overwrite_force"
-		suite.destinationPluginTestMigrate(ctx, mem, t, newPlugin, logger, destSpec, tests.MigrateStrategyOverwrite)
+		suite.destinationPluginTestMigrate(ctx, t, newPlugin, logger, destSpec, tests.MigrateStrategyOverwrite)
 	})
 
 	t.Run("TestWriteAppend", func(t *testing.T) {
 		t.Helper()
 		if suite.tests.SkipAppend {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.Name = "test_write_append"
 		p := newPlugin()
-		if err := suite.destinationPluginTestWriteAppend(ctx, mem, p, logger, destSpec); err != nil {
+		if err := suite.destinationPluginTestWriteAppend(ctx, p, logger, destSpec); err != nil {
 			t.Fatal(err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -189,27 +178,23 @@ func PluginTestSuiteRunner(t *testing.T, newPlugin NewPluginFunc, destSpec specs
 
 	t.Run("TestMigrateAppend", func(t *testing.T) {
 		t.Helper()
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		if suite.tests.SkipMigrateAppend {
 			t.Skip("skipping " + t.Name())
 		}
 		destSpec.WriteMode = specs.WriteModeAppend
 		destSpec.Name = "test_migrate_append"
-		suite.destinationPluginTestMigrate(ctx, mem, t, newPlugin, logger, destSpec, tests.MigrateStrategyAppend)
+		suite.destinationPluginTestMigrate(ctx, t, newPlugin, logger, destSpec, tests.MigrateStrategyAppend)
 	})
 
 	t.Run("TestMigrateAppendForce", func(t *testing.T) {
 		t.Helper()
 		if suite.tests.SkipMigrateAppendForce {
 			t.Skip("skipping " + t.Name())
 		}
-		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
-		defer mem.AssertSize(t, 0)
 		destSpec.WriteMode = specs.WriteModeAppend
 		destSpec.MigrateMode = specs.MigrateModeForced
 		destSpec.Name = "test_migrate_append_force"
-		suite.destinationPluginTestMigrate(ctx, mem, t, newPlugin, logger, destSpec, tests.MigrateStrategyAppend)
+		suite.destinationPluginTestMigrate(ctx, t, newPlugin, logger, destSpec, tests.MigrateStrategyAppend)
 	})
 }
 
diff --git a/plugins/destination/plugin_testing_migrate.go b/plugins/destination/plugin_testing_migrate.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/array"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v2/schema"
 	"github.com/cloudquery/plugin-sdk/v2/specs"
 	"github.com/cloudquery/plugin-sdk/v2/testdata"
@@ -22,7 +21,7 @@ func tableUUIDSuffix() string {
 	return strings.ReplaceAll(uuid.NewString(), "-", "_")
 }
 
-func testMigration(ctx context.Context, mem memory.Allocator, _ *testing.T, p *Plugin, logger zerolog.Logger, spec specs.Destination, target *arrow.Schema, source *arrow.Schema, mode specs.MigrateMode) error {
+func testMigration(ctx context.Context, _ *testing.T, p *Plugin, logger zerolog.Logger, spec specs.Destination, target *arrow.Schema, source *arrow.Schema, mode specs.MigrateMode) error {
 	if err := p.Init(ctx, logger, spec); err != nil {
 		return fmt.Errorf("failed to init plugin: %w", err)
 	}
@@ -41,19 +40,15 @@ func testMigration(ctx context.Context, mem memory.Allocator, _ *testing.T, p *P
 		SyncTime:   syncTime,
 		MaxRows:    1,
 	}
-	resource1 := testdata.GenTestData(mem, source, opts)[0]
-	resource1.Retain()
-	defer resource1.Release()
+	resource1 := testdata.GenTestData(source, opts)[0]
 	if err := p.writeOne(ctx, sourceSpec, syncTime, resource1); err != nil {
 		return fmt.Errorf("failed to write one: %w", err)
 	}
 
 	if err := p.Migrate(ctx, []*arrow.Schema{target}); err != nil {
 		return fmt.Errorf("failed to migrate existing table: %w", err)
 	}
-	resource2 := testdata.GenTestData(mem, target, opts)[0]
-	resource2.Retain()
-	defer resource2.Release()
+	resource2 := testdata.GenTestData(target, opts)[0]
 	if err := p.writeOne(ctx, sourceSpec, syncTime, resource2); err != nil {
 		return fmt.Errorf("failed to write one after migration: %w", err)
 	}
@@ -85,7 +80,6 @@ func testMigration(ctx context.Context, mem memory.Allocator, _ *testing.T, p *P
 
 func (*PluginTestSuite) destinationPluginTestMigrate(
 	ctx context.Context,
-	mem memory.Allocator,
 	t *testing.T,
 	newPlugin NewPluginFunc,
 	logger zerolog.Logger,
@@ -117,7 +111,7 @@ func (*PluginTestSuite) destinationPluginTestMigrate(
 		}, &md)
 
 		p := newPlugin()
-		if err := testMigration(ctx, mem, t, p, logger, spec, target, source, strategy.AddColumn); err != nil {
+		if err := testMigration(ctx, t, p, logger, spec, target, source, strategy.AddColumn); err != nil {
 			t.Fatalf("failed to migrate %s: %v", tableName, err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -147,7 +141,7 @@ func (*PluginTestSuite) destinationPluginTestMigrate(
 			{Name: "bool", Type: arrow.FixedWidthTypes.Boolean},
 		}, &md)
 		p := newPlugin()
-		if err := testMigration(ctx, mem, t, p, logger, spec, target, source, strategy.AddColumnNotNull); err != nil {
+		if err := testMigration(ctx, t, p, logger, spec, target, source, strategy.AddColumnNotNull); err != nil {
 			t.Fatalf("failed to migrate add_column_not_null: %v", err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -177,7 +171,7 @@ func (*PluginTestSuite) destinationPluginTestMigrate(
 		}, &md)
 
 		p := newPlugin()
-		if err := testMigration(ctx, mem, t, p, logger, spec, target, source, strategy.RemoveColumn); err != nil {
+		if err := testMigration(ctx, t, p, logger, spec, target, source, strategy.RemoveColumn); err != nil {
 			t.Fatalf("failed to migrate remove_column: %v", err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -207,7 +201,7 @@ func (*PluginTestSuite) destinationPluginTestMigrate(
 		}, &md)
 
 		p := newPlugin()
-		if err := testMigration(ctx, mem, t, p, logger, spec, target, source, strategy.RemoveColumnNotNull); err != nil {
+		if err := testMigration(ctx, t, p, logger, spec, target, source, strategy.RemoveColumnNotNull); err != nil {
 			t.Fatalf("failed to migrate remove_column_not_null: %v", err)
 		}
 		if err := p.Close(ctx); err != nil {
@@ -238,7 +232,7 @@ func (*PluginTestSuite) destinationPluginTestMigrate(
 		}, &md)
 
 		p := newPlugin()
-		if err := testMigration(ctx, mem, t, p, logger, spec, target, source, strategy.ChangeColumn); err != nil {
+		if err := testMigration(ctx, t, p, logger, spec, target, source, strategy.ChangeColumn); err != nil {
 			t.Fatalf("failed to migrate change_column: %v", err)
 		}
 		if err := p.Close(ctx); err != nil {
diff --git a/plugins/destination/plugin_testing_overwrite.go b/plugins/destination/plugin_testing_overwrite.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/array"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v2/schema"
 	"github.com/cloudquery/plugin-sdk/v2/specs"
 	"github.com/cloudquery/plugin-sdk/v2/testdata"
@@ -16,7 +15,7 @@ import (
 	"github.com/rs/zerolog"
 )
 
-func (*PluginTestSuite) destinationPluginTestWriteOverwrite(ctx context.Context, mem memory.Allocator, p *Plugin, logger zerolog.Logger, spec specs.Destination) error {
+func (*PluginTestSuite) destinationPluginTestWriteOverwrite(ctx context.Context, p *Plugin, logger zerolog.Logger, spec specs.Destination) error {
 	spec.WriteMode = specs.WriteModeOverwrite
 	if err := p.Init(ctx, logger, spec); err != nil {
 		return fmt.Errorf("failed to init plugin: %w", err)
@@ -41,15 +40,7 @@ func (*PluginTestSuite) destinationPluginTestWriteOverwrite(ctx context.Context,
 		SyncTime:   syncTime,
 		MaxRows:    2,
 	}
-	resources := testdata.GenTestData(mem, schema.CQSchemaToArrow(table), opts)
-	for _, r := range resources {
-		r.Retain()
-	}
-	defer func() {
-		for _, r := range resources {
-			r.Release()
-		}
-	}()
+	resources := testdata.GenTestData(schema.CQSchemaToArrow(table), opts)
 	if err := p.writeAll(ctx, sourceSpec, syncTime, resources); err != nil {
 		return fmt.Errorf("failed to write all: %w", err)
 	}
@@ -85,9 +76,7 @@ func (*PluginTestSuite) destinationPluginTestWriteOverwrite(ctx context.Context,
 		MaxRows:    1,
 		StableUUID: *u,
 	}
-	updatedResource := testdata.GenTestData(mem, schema.CQSchemaToArrow(table), opts)[0]
-	updatedResource.Retain()
-	defer updatedResource.Release()
+	updatedResource := testdata.GenTestData(schema.CQSchemaToArrow(table), opts)[0]
 	// write second time
 	if err := p.writeOne(ctx, sourceSpec, secondSyncTime, updatedResource); err != nil {
 		return fmt.Errorf("failed to write one second time: %w", err)
diff --git a/plugins/destination/plugin_testing_overwrite_delete_stale.go b/plugins/destination/plugin_testing_overwrite_delete_stale.go
@@ -7,15 +7,14 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/array"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/cloudquery/plugin-sdk/v2/specs"
 	"github.com/cloudquery/plugin-sdk/v2/testdata"
 	"github.com/cloudquery/plugin-sdk/v2/types"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
 )
 
-func (*PluginTestSuite) destinationPluginTestWriteOverwriteDeleteStale(ctx context.Context, mem memory.Allocator, p *Plugin, logger zerolog.Logger, spec specs.Destination) error {
+func (*PluginTestSuite) destinationPluginTestWriteOverwriteDeleteStale(ctx context.Context, p *Plugin, logger zerolog.Logger, spec specs.Destination) error {
 	spec.WriteMode = specs.WriteModeOverwriteDeleteStale
 	if err := p.Init(ctx, logger, spec); err != nil {
 		return fmt.Errorf("failed to init plugin: %w", err)
@@ -44,18 +43,10 @@ func (*PluginTestSuite) destinationPluginTestWriteOverwriteDeleteStale(ctx conte
 		SyncTime:   syncTime,
 		MaxRows:    2,
 	}
-	resources := testdata.GenTestData(mem, table.ToArrowSchema(), opts)
-	incResources := testdata.GenTestData(mem, incTable.ToArrowSchema(), opts)
+	resources := testdata.GenTestData(table.ToArrowSchema(), opts)
+	incResources := testdata.GenTestData(incTable.ToArrowSchema(), opts)
 	allResources := resources
 	allResources = append(allResources, incResources...)
-	for _, r := range allResources {
-		r.Retain()
-	}
-	defer func() {
-		for _, r := range allResources {
-			r.Release()
-		}
-	}()
 	if err := p.writeAll(ctx, sourceSpec, syncTime, allResources); err != nil {
 		return fmt.Errorf("failed to write all: %w", err)
 	}
@@ -98,9 +89,7 @@ func (*PluginTestSuite) destinationPluginTestWriteOverwriteDeleteStale(ctx conte
 		StableUUID: *u,
 		MaxRows:    1,
 	}
-	updatedResources := testdata.GenTestData(mem, table.ToArrowSchema(), opts)[0]
-	updatedResources.Retain()
-	defer updatedResources.Release()
+	updatedResources := testdata.GenTestData(table.ToArrowSchema(), opts)[0]
 
 	if err := p.writeOne(ctx, sourceSpec, secondSyncTime, updatedResources); err != nil {
 		return fmt.Errorf("failed to write one second time: %w", err)
diff --git a/plugins/destination/plugin_testing_write_append.go b/plugins/destination/plugin_testing_write_append.go
diff --git a/testdata/testdata.go b/testdata/testdata.go

Original file line number	Diff line number	Diff line change
`@@ -86,10 +86,8 @@ func (c client) overwrite(table arrow.Schema, data arrow.Record) {`
`86`	`86`	`}`
`87`	`87`	`}`
`88`	`88`	`if found {`
`89`		`- tmp := c.memoryDB[tableName][i]`
`90`	`89`	`c.memoryDB[tableName] = append(c.memoryDB[tableName][:i], c.memoryDB[tableName][i+1:]...)`
`91`	`90`	`c.memoryDB[tableName] = append(c.memoryDB[tableName], data)`
`92`		`- tmp.Release()`
`93`	`91`	`return`
`94`	`92`	`}`
`95`	`93`	`}`
`@@ -110,11 +108,6 @@ func (c *client) Migrate(_ context.Context, tables schema.Schemas) error {`
`110`	`108`	`if changes == nil {`
`111`	`109`	`continue`
`112`	`110`	`}`
`113`		`- for _, t := range c.memoryDB {`
`114`		`- for _, row := range t {`
`115`		`- row.Release()`
`116`		`- }`
`117`		`- }`
`118`	`111`	`c.memoryDB[tableName] = make([]arrow.Record, 0)`
`119`	`112`	`c.tables[tableName] = table`
`120`	`113`	`}`
`@@ -204,11 +197,6 @@ func (*client) Metrics() destination.Metrics {`
`204`	`197`	`}`
`205`	`198`
`206`	`199`	`func (c *client) Close(context.Context) error {`
`207`		`- for _, table := range c.memoryDB {`
`208`		`- for _, row := range table {`
`209`		`- row.Release()`
`210`		`- }`
`211`		`- }`
`212`	`200`	`c.memoryDB = nil`
`213`	`201`	`return nil`
`214`	`202`	`}`
`@@ -230,8 +218,6 @@ func (c client) deleteStaleTable(_ context.Context, table arrow.Schema, source`
`230`	`218`	`rowSyncTime := row.Column(syncColIndex).(*array.Timestamp).Value(0).ToTime(arrow.Microsecond).UTC()`
`231`	`219`	`if !rowSyncTime.Before(syncTime) {`
`232`	`220`	`filteredTable = append(filteredTable, c.memoryDB[tableName][i])`
`233`		`- } else {`
`234`		`- c.memoryDB[tableName][i].Release()`
`235`	`221`	`}`
`236`	`222`	`}`
`237`	`223`	`}`