This is an automated email from the ASF dual-hosted git repository.
chaokunyang pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fory.git
The following commit(s) were added to refs/heads/main by this push:
new 697295af4 fix: fix go thread safe Pool race condition and remove
python lazy_import (#4027)
697295af4 is described below
commit 697295af47c121979c7d6910f32ff0cff3b72651
Author: Shawn Yang <[email protected]>
AuthorDate: Tue Sep 8 00:40:42 2026 +0800
fix: fix go thread safe Pool race condition and remove python lazy_import
(#4027)
## Why?
## What does this PR do?
## Related issues
## AI Contribution Checklist
- [ ] Substantial AI assistance was used in this PR: `yes` / `no`
- [ ] If `yes`, I included a completed [AI Contribution
Checklist](https://github.com/apache/fory/blob/main/AI_POLICY.md#9-contributor-checklist-for-ai-assisted-prs)
in this PR description and the required `AI Usage Disclosure`.
- [ ] If `yes`, my PR description includes the required `ai_review`
summary and screenshot evidence or equivalent persisted links of the
final clean AI review results from both fresh reviewers described in
`AI_POLICY.md`, the Fory-guided reviewer and the independent general
reviewer, on the current PR diff or current HEAD after the latest code
changes.
## Does this PR introduce any user-facing change?
- [ ] Does this PR introduce any public API change?
- [ ] Does this PR introduce any binary protocol compatibility change?
## Benchmark
---
.agents/languages/go.md | 9 +
ci/run_ci.sh | 2 +-
ci/tasks/go.py | 2 +-
docs/object-serialization/go/configuration.md | 2 +-
docs/object-serialization/go/thread-safety.md | 75 ++++----
go/fory/threadsafe/fory.go | 128 +++++++++++--
go/fory/threadsafe/fory_test.go | 249 ++++++++++++++++++++++++++
python/pyfory/format/tests/test_vectorized.py | 5 +-
python/pyfory/tests/record.py | 7 +-
python/pyfory/tests/test_buffer.py | 6 +-
python/pyfory/tests/test_cross_language.py | 9 +-
python/pyfory/tests/test_serializer.py | 4 +-
python/pyfory/utils.py | 63 -------
13 files changed, 423 insertions(+), 138 deletions(-)
diff --git a/.agents/languages/go.md b/.agents/languages/go.md
index 5192fc3e7..ea550e49c 100644
--- a/.agents/languages/go.md
+++ b/.agents/languages/go.md
@@ -7,6 +7,15 @@ Load this file when changing `go/fory/` or Go xlang behavior.
- Run Go commands from within `go/fory/`.
- Changes under `go/` must pass formatting and tests.
- The Go implementation focuses on fast serializers.
+- Keep `threadsafe.Fory` backed by `sync.Pool`. Direct struct and enum
registration
+ by ID or name must initialize every instance, including replacements after
GC.
+ Follow Java `ThreadLocalFory`'s registration callback semantics without
copying
+ its thread-local storage or replacing the Go pool with a bounded pool.
+ Configure the first actual instance before pooling it; registration cannot
+ accumulate on arbitrary `sync.Pool.Get` results. Freeze successful callbacks
+ before the first root, apply them only when creating additional instances,
and
+ leave pool hits free of registration replay. Preserve direct registration
+ examples rather than forcing callers to factories.
- Go `ReadContext` intentionally defers codec errors to existing `HasError` or
`CheckError`
boundaries. After an error, work may continue only while it remains panic-
and bounds-safe and
cannot cause disproportionate work or allocation, publish state that
survives root cleanup, or
diff --git a/ci/run_ci.sh b/ci/run_ci.sh
index ebdab570b..a56dc8337 100755
--- a/ci/run_ci.sh
+++ b/ci/run_ci.sh
@@ -422,7 +422,7 @@ case $1 in
go)
echo "Executing fory go tests for go"
cd "$ROOT/go/fory"
- go test -v ./...
+ go test -race -v ./...
echo "Executing fory go tests succeeds"
;;
format)
diff --git a/ci/tasks/go.py b/ci/tasks/go.py
index 3df17f81f..e4e960dc1 100644
--- a/ci/tasks/go.py
+++ b/ci/tasks/go.py
@@ -23,5 +23,5 @@ def run():
"""Run Go CI tasks."""
logging.info("Executing fory go tests")
common.cd_project_subdir("go/fory")
- common.exec_cmd("go test -v")
+ common.exec_cmd("go test -race -v ./...")
logging.info("Executing fory go tests succeeds")
diff --git a/docs/object-serialization/go/configuration.md
b/docs/object-serialization/go/configuration.md
index 8f59db06e..50850328d 100644
--- a/docs/object-serialization/go/configuration.md
+++ b/docs/object-serialization/go/configuration.md
@@ -241,7 +241,7 @@ go func() {
The thread-safe wrapper:
-- Uses `sync.Pool` internally for efficient instance reuse
+- Creates instances as needed and reuses them across goroutines
- Automatically copies serialized data before returning
- Accepts the same configuration options as `fory.New()`
diff --git a/docs/object-serialization/go/thread-safety.md
b/docs/object-serialization/go/thread-safety.md
index 89b536766..a9ed08dfc 100644
--- a/docs/object-serialization/go/thread-safety.md
+++ b/docs/object-serialization/go/thread-safety.md
@@ -66,32 +66,11 @@ go func() {
}()
```
-### How It Works
-
-The thread-safe wrapper uses `sync.Pool`:
-
-1. **Acquire**: Gets a Fory instance from the pool
-2. **Use**: Performs serialization/deserialization
-3. **Copy**: Copies result data (buffer will be reused)
-4. **Release**: Returns instance to pool
-
-```go
-// Simplified implementation
-func (f *Fory) Serialize(v any) ([]byte, error) {
- fory := f.pool.Get().(*fory.Fory)
- defer f.pool.Put(fory)
-
- data, err := fory.Serialize(v)
- if err != nil {
- return nil, err
- }
-
- // Copy because underlying buffer will be reused
- result := make([]byte, len(data))
- copy(result, data)
- return result, nil
-}
-```
+The wrapper creates instances as needed and reuses them across goroutines.
+Each operation exclusively borrows one instance and returns it afterward.
+Registered types remain available when garbage collection reclaims cached
+instances. Serialized output is copied before returning, so callers can retain
+it safely.
### API
@@ -114,32 +93,48 @@ err = threadsafe.Unmarshal(data, &target)
## Type Registration
-Type registration should be done before concurrent use:
+Register all types before the first serialization or deserialization
operation. The
+first operation permanently freezes the wrapper's registrations, even if that
+operation fails. A later registration attempt returns an error.
```go
f := threadsafe.New()
-// Register types BEFORE concurrent access
-f.RegisterStruct(User{}, 1)
-f.RegisterStruct(Order{}, 2)
+if err := f.RegisterStruct(User{}, 1); err != nil {
+ panic(err)
+}
+if err := f.RegisterStruct(Order{}, 2); err != nil {
+ panic(err)
+}
-// Now safe to use concurrently
+// All concurrent operations use the registered types.
go func() {
- f.Serialize(&User{ID: 1})
+ data, err := f.Serialize(&User{ID: 1})
+ // Use data and handle err.
+ _, _ = data, err
}()
```
-### Thread-Safe Registration
+`RegisterStructByName`, `RegisterEnum`, and `RegisterEnumByName` are also
+available directly on the wrapper. Every registered type is available to all
+concurrent operations and remains registered across garbage collections.
-The thread-safe wrapper handles registration safely:
+For custom per-instance initialization, use `NewWithFactory`:
```go
-// Safe: Registration is synchronized
-f := threadsafe.New()
-f.RegisterStruct(User{}, 1) // Thread-safe
+f := threadsafe.NewWithFactory(func() *fory.Fory {
+ inner := fory.New()
+ if err := inner.RegisterExtension(CustomType{}, 100,
newCustomSerializer()); err != nil {
+ panic(err)
+ }
+ return inner
+})
```
-However, for best performance, register all types at startup before concurrent
use.
+The factory may be called concurrently as additional instances are needed. It
+must return a fresh, identically configured instance on every call, with
+registrations completed before returning. Create stateful custom serializers
+separately for each instance.
## Zero-Copy Considerations
@@ -331,11 +326,11 @@ go func() {
}()
```
-**Fix**: Register all types before concurrent use.
+**Fix**: Register all types before the first serialization or deserialization.
## Best Practices
-1. **Register types at startup**: Before any concurrent operations
+1. **Register types at startup**: Before the first serialization or
deserialization
2. **Clone data if keeping references**: With non-thread-safe instance
3. **Use per-worker instances for hot paths**: Eliminates pool contention
4. **Profile before optimizing**: Thread-safe overhead may be negligible
diff --git a/go/fory/threadsafe/fory.go b/go/fory/threadsafe/fory.go
index 4afdfa0de..8593d4b1d 100644
--- a/go/fory/threadsafe/fory.go
+++ b/go/fory/threadsafe/fory.go
@@ -15,19 +15,27 @@
// specific language governing permissions and limitations
// under the License.
-// Package threadsafe provides a thread-safe wrapper around Fory using
sync.Pool.
+// Package threadsafe provides a thread-safe wrapper around Fory.
package threadsafe
import (
+ "fmt"
+ "reflect"
"sync"
+ "sync/atomic"
"github.com/apache/fory/go/fory"
)
// Fory is a thread-safe wrapper around fory.Fory using sync.Pool.
-// It provides the same API as fory.Fory but is safe for concurrent use.
+// Struct and enum registration use the same API as fory.Fory and must finish
+// before the first serialization or deserialization operation.
type Fory struct {
- pool sync.Pool
+ pool sync.Pool
+ registrationMu sync.Mutex
+ started atomic.Bool
+ first *fory.Fory
+ callbacks []func(*fory.Fory) error
}
// New creates a new thread-safe Fory instance.
@@ -37,28 +45,54 @@ func New(opts ...fory.Option) *Fory {
})
}
-// NewWithFactory creates a new thread-safe Fory instance using a custom
factory.
+// NewWithFactory creates a thread-safe Fory using a custom factory.
+// The factory must return a fresh, identically configured Fory instance on
every
+// call, with any custom registrations completed before returning. It may be
+// called concurrently when additional instances are needed.
func NewWithFactory(factory func() *fory.Fory) *Fory {
if factory == nil {
panic("threadsafe.NewWithFactory requires a non-nil factory")
}
f := &Fory{}
- f.pool = sync.Pool{
- New: func() any {
- inner := factory()
- if inner == nil {
- panic("threadsafe.NewWithFactory factory
returned nil")
+ f.pool.New = func() any {
+ inner := factory()
+ if inner == nil {
+ panic("threadsafe.NewWithFactory factory returned nil")
+ }
+ // Like Java ThreadLocalFory's factoryCallback, registrations
initialize
+ // every new instance, including replacements for entries
discarded by GC.
+ // Setup holds registrationMu; root operations freeze callbacks
before Get.
+ for _, callback := range f.callbacks {
+ if err := callback(inner); err != nil {
+ panic(fmt.Errorf("threadsafe factory
registration failed: %w", err))
}
- return inner
- },
+ }
+ return inner
}
return f
}
func (f *Fory) acquire() *fory.Fory {
+ if !f.started.Load() {
+ if inner := f.freezeRegistrations(); inner != nil {
+ return inner
+ }
+ }
return f.pool.Get().(*fory.Fory)
}
+//go:noinline
+func (f *Fory) freezeRegistrations() *fory.Fory {
+ f.registrationMu.Lock()
+ defer f.registrationMu.Unlock()
+ // The first root directly borrows the instance configured during setup.
+ // Freeze even if that root fails, and let sync.Pool own its reuse
afterward.
+ inner := f.first
+ f.first = nil
+ f.started.Store(true)
+ return inner
+}
+
func (f *Fory) release(inner *fory.Fory) {
inner.Reset()
f.pool.Put(inner)
@@ -90,11 +124,75 @@ func (f *Fory) Deserialize(data []byte, v any) error {
return inner.Deserialize(data, v)
}
-// RegisterStructByName registers a struct type by name for cross-language
serialization.
+func (f *Fory) registerCallback(registration func(*fory.Fory) error) error {
+ f.registrationMu.Lock()
+ defer f.registrationMu.Unlock()
+ if f.started.Load() {
+ return fmt.Errorf("types must be registered before the first
serialization or deserialization")
+ }
+ // Registration cannot accumulate on arbitrary pool.Get results: even
+ // serialized callers can borrow different entries, and GC can discard
them.
+ // Configure the first real instance before it enters the pool instead.
+ if f.first == nil {
+ f.first = f.pool.New().(*fory.Fory)
+ }
+ if err := registration(f.first); err != nil {
+ return err
+ }
+ f.callbacks = append(f.callbacks, registration)
+ return nil
+}
+
+// Keep only type metadata in callbacks, rather than retaining caller objects.
+func registrationType(type_ any) reflect.Type {
+ if typ, ok := type_.(reflect.Type); ok {
+ return typ
+ }
+ typ := reflect.TypeOf(type_)
+ if typ != nil && typ.Kind() == reflect.Ptr {
+ typ = typ.Elem()
+ }
+ return typ
+}
+
+// RegisterStruct registers a struct type with a numeric ID in every pooled
instance.
+// Registration must complete before the first serialization or
deserialization,
+// including a failed operation. Arguments follow fory.Fory.RegisterStruct.
+func (f *Fory) RegisterStruct(type_ any, typeID uint32) error {
+ typ := registrationType(type_)
+ return f.registerCallback(func(inner *fory.Fory) error {
+ return inner.RegisterStruct(typ, typeID)
+ })
+}
+
+// RegisterStructByName registers a struct type by name in every pooled
instance.
+// Registration must complete before the first serialization or
deserialization,
+// including a failed operation. Arguments follow
fory.Fory.RegisterStructByName.
func (f *Fory) RegisterStructByName(type_ any, name string) error {
- inner := f.acquire()
- defer f.release(inner)
- return inner.RegisterStructByName(type_, name)
+ typ := registrationType(type_)
+ return f.registerCallback(func(inner *fory.Fory) error {
+ return inner.RegisterStructByName(typ, name)
+ })
+}
+
+// RegisterEnum registers an enum type with a numeric ID in every pooled
instance.
+// Registration must complete before the first serialization or
deserialization,
+// including a failed operation. Arguments follow fory.Fory.RegisterEnum.
+func (f *Fory) RegisterEnum(type_ any, typeID uint32) error {
+ typ := registrationType(type_)
+ return f.registerCallback(func(inner *fory.Fory) error {
+ return inner.RegisterEnum(typ, typeID)
+ })
+}
+
+// RegisterEnumByName registers an enum type by name in every pooled instance.
+// Registration must complete before the first serialization or
deserialization,
+// including a failed operation. Arguments follow fory.Fory.RegisterEnumByName.
+func (f *Fory) RegisterEnumByName(type_ any, name string) error {
+ typ := registrationType(type_)
+ return f.registerCallback(func(inner *fory.Fory) error {
+ return inner.RegisterEnumByName(typ, name)
+ })
}
// ============================================================================
diff --git a/go/fory/threadsafe/fory_test.go b/go/fory/threadsafe/fory_test.go
index 37b4a0ecf..b7272d325 100644
--- a/go/fory/threadsafe/fory_test.go
+++ b/go/fory/threadsafe/fory_test.go
@@ -18,12 +18,260 @@
package threadsafe
import (
+ "reflect"
+ "runtime"
+ "sync"
+ "sync/atomic"
"testing"
"github.com/apache/fory/go/fory"
"github.com/stretchr/testify/require"
)
+func TestRegistrationAfterGC(t *testing.T) {
+ type Item struct{ Value int32 }
+ for _, byName := range []bool{false, true} {
+ name := "ID"
+ if byName {
+ name = "Name"
+ }
+ t.Run(name, func(t *testing.T) {
+ var created atomic.Int32
+ f := NewWithFactory(func() *fory.Fory {
+ created.Add(1)
+ return fory.New()
+ })
+ if byName {
+ require.NoError(t,
f.RegisterStructByName(Item{}, "threadsafe.Item"))
+ } else {
+ require.NoError(t, f.RegisterStruct(Item{}, 1))
+ }
+ runtime.GC()
+ runtime.GC()
+ value := Item{Value: 42}
+ data, err := f.Serialize(&value)
+ require.NoError(t, err)
+ // The configured instance serves the first root, even
after setup GC.
+ require.Equal(t, int32(1), created.Load())
+
+ // Two collections discard the sync.Pool primary and
victim caches.
+ runtime.GC()
+ runtime.GC()
+ var result Item
+ require.NoError(t, f.Deserialize(data, &result))
+ require.Equal(t, value, result)
+ require.Greater(t, created.Load(), int32(1))
+
+ before := created.Load()
+ runtime.GC()
+ runtime.GC()
+ data, err = Serialize(f, &value)
+ require.NoError(t, err)
+ require.Greater(t, created.Load(), before)
+ require.NoError(t, Deserialize(f, data, &result))
+ require.Equal(t, value, result)
+ })
+ }
+}
+
+func TestRegistrationInstances(t *testing.T) {
+ type Order struct{ ID int64 }
+ type User struct {
+ ID int32
+ Order *Order
+ }
+ type NamedItem struct{ Name string }
+ type FactoryItem struct{ Name string }
+ type State int32
+ type Color int32
+ var factoryCalls atomic.Int32
+ var firstCreated *fory.Fory
+ f := NewWithFactory(func() *fory.Fory {
+ inner := fory.New()
+ if factoryCalls.Add(1) == 1 {
+ firstCreated = inner
+ }
+ if err := inner.RegisterStruct(FactoryItem{}, 100); err != nil {
+ panic(err)
+ }
+ return inner
+ })
+ require.Zero(t, factoryCalls.Load())
+ require.Error(t, f.RegisterStructByName(int32(0), "threadsafe.Invalid"))
+ require.NoError(t, f.RegisterStruct(&User{}, 1))
+ runtime.GC()
+ runtime.GC()
+ require.NoError(t, f.RegisterStruct(reflect.TypeOf(User{}), 1))
+ require.Error(t, f.RegisterStruct(Order{}, 1))
+ require.Error(t, f.RegisterStructByName(User{}, "threadsafe.Duplicate"))
+ require.NoError(t, f.RegisterStruct(Order{}, 2))
+ require.NoError(t, f.RegisterStructByName(NamedItem{},
"threadsafe.NamedItem"))
+ require.NoError(t, f.RegisterEnum(State(0), 3))
+ require.NoError(t, f.RegisterEnumByName(reflect.TypeOf(Color(0)),
"threadsafe.Color"))
+ // Ordinary registration does not rebuild the factory's instances each
time.
+ require.Equal(t, int32(1), factoryCalls.Load())
+
+ // Hold both borrows to force an additional, independently configured
instance.
+ first := f.acquire()
+ second := f.acquire()
+ require.Same(t, firstCreated, first)
+ require.NotSame(t, first, second)
+ require.Equal(t, int32(2), factoryCalls.Load())
+ for _, inner := range []*fory.Fory{first, second} {
+ for _, value := range []any{&User{42, &Order{7}}, &Order{7},
&NamedItem{"name"}, &FactoryItem{"factory"}, State(1), Color(2)} {
+ data, err := inner.Serialize(value)
+ require.NoError(t, err)
+ var result any
+ require.NoError(t, inner.Deserialize(data, &result))
+ require.Equal(t, value, result)
+ }
+ }
+ f.release(first)
+ f.release(second)
+
+ var workers sync.WaitGroup
+ for range 8 {
+ workers.Add(1)
+ go func() {
+ defer workers.Done()
+ value := User{42, &Order{7}}
+ for range 20 {
+ data, err := Serialize(f, &value)
+ if err != nil {
+ t.Error(err)
+ return
+ }
+ var result User
+ if err := Deserialize(f, data, &result); err !=
nil {
+ t.Error(err)
+ return
+ }
+ if result.ID != value.ID || result.Order == nil
|| *result.Order != *value.Order {
+ t.Errorf("got %v, want %v", result,
value)
+ return
+ }
+ }
+ }()
+ }
+ workers.Wait()
+}
+
+func TestConcurrentRegistration(t *testing.T) {
+ type User struct{ ID int32 }
+ type Order struct{ ID int64 }
+ type State int32
+ type Color int32
+ var created atomic.Int32
+ f := NewWithFactory(func() *fory.Fory {
+ created.Add(1)
+ return fory.New()
+ })
+ register := []func() error{
+ func() error { return f.RegisterStruct(User{}, 1) },
+ func() error { return f.RegisterStructByName(Order{},
"threadsafe.Order") },
+ func() error { return f.RegisterEnum(State(0), 2) },
+ func() error { return f.RegisterEnumByName(Color(0),
"threadsafe.Color") },
+ }
+ var workers sync.WaitGroup
+ for _, call := range register {
+ workers.Add(1)
+ go func() {
+ defer workers.Done()
+ if err := call(); err != nil {
+ t.Error(err)
+ }
+ }()
+ }
+ workers.Wait()
+ require.Equal(t, int32(1), created.Load())
+ first := f.acquire()
+ second := f.acquire()
+ for _, inner := range []*fory.Fory{first, second} {
+ for _, value := range []any{&User{42}, &Order{7}, State(1),
Color(2)} {
+ data, err := inner.Serialize(value)
+ require.NoError(t, err)
+ var result any
+ require.NoError(t, inner.Deserialize(data, &result))
+ require.Equal(t, value, result)
+ }
+ }
+ f.release(first)
+ f.release(second)
+}
+
+func TestRegistrationAtFirstRoot(t *testing.T) {
+ type User struct{ ID int32 }
+ type Order struct{ ID int64 }
+ for range 20 {
+ f := New()
+ require.NoError(t, f.RegisterStruct(User{}, 1))
+ start := make(chan struct{})
+ registered := make(chan error, 1)
+ serialized := make(chan error, 1)
+ go func() {
+ <-start
+ registered <- f.RegisterStruct(Order{}, 2)
+ }()
+ go func() {
+ <-start
+ _, err := f.Serialize(&User{42})
+ serialized <- err
+ }()
+ close(start)
+ registrationErr := <-registered
+ require.NoError(t, <-serialized)
+ first := f.acquire()
+ second := f.acquire()
+ for _, inner := range []*fory.Fory{first, second} {
+ _, err := inner.Serialize(&Order{7})
+ if registrationErr == nil {
+ require.NoError(t, err)
+ } else {
+ require.Error(t, err)
+ }
+ }
+ f.release(first)
+ f.release(second)
+ }
+}
+
+func TestRegistrationFreeze(t *testing.T) {
+ type Item struct{ Value int32 }
+ value := int32(42)
+ data, err := fory.New().Serialize(value)
+ require.NoError(t, err)
+ tests := []struct {
+ name string
+ fail bool
+ call func(*Fory) error
+ }{
+ {"Serialize", false, func(f *Fory) error { _, err :=
f.Serialize(value); return err }},
+ {"SerializeError", true, func(f *Fory) error { _, err :=
f.Serialize(Item{}); return err }},
+ {"GenericSerialize", false, func(f *Fory) error { _, err :=
Serialize(f, &value); return err }},
+ {"GenericSerializeError", true, func(f *Fory) error { _, err :=
Serialize(f, &Item{}); return err }},
+ {"Deserialize", false, func(f *Fory) error { return
f.Deserialize(data, new(int32)) }},
+ {"DeserializeError", true, func(f *Fory) error { return
f.Deserialize(nil, new(int32)) }},
+ {"GenericDeserialize", false, func(f *Fory) error { return
Deserialize(f, data, new(int32)) }},
+ {"GenericDeserializeError", true, func(f *Fory) error { return
Deserialize(f, nil, new(int32)) }},
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ f := New()
+ err := test.call(f)
+ if test.fail {
+ require.Error(t, err)
+ } else {
+ require.NoError(t, err)
+ }
+ require.Error(t, f.RegisterStruct(Item{}, 1))
+ require.Error(t, f.RegisterStructByName(Item{},
"threadsafe.Item"))
+ type State int32
+ require.Error(t, f.RegisterEnum(State(0), 2))
+ require.Error(t, f.RegisterEnumByName(State(0),
"threadsafe.State"))
+ })
+ }
+}
+
// TestFory tests the thread-safe Fory wrapper
func TestFory(t *testing.T) {
f := New(fory.WithXlang(false), fory.WithRefTracking(true),
fory.WithCompatible(false))
@@ -140,6 +388,7 @@ func TestDeserialize(t *testing.T) {
})
t.Run("Slice", func(t *testing.T) {
+ f := New(fory.WithXlang(false), fory.WithRefTracking(true),
fory.WithCompatible(false))
// Serialize a struct containing the slice since *[]T is not
supported
type SliceWrapper struct {
Items []int32
diff --git a/python/pyfory/format/tests/test_vectorized.py
b/python/pyfory/format/tests/test_vectorized.py
index 2af57e555..ba2bf49a1 100644
--- a/python/pyfory/format/tests/test_vectorized.py
+++ b/python/pyfory/format/tests/test_vectorized.py
@@ -17,8 +17,7 @@
import pyfory as fory
-from pyfory.tests.core import require_pyarrow
-from pyfory.utils import lazy_import
+from pyfory.tests.core import pa, require_pyarrow
from pyfory.format import (
schema,
field,
@@ -31,8 +30,6 @@ from pyfory.format import (
to_arrow_schema,
)
-pa = lazy_import("pyarrow")
-
@require_pyarrow
def test_vectorized():
diff --git a/python/pyfory/tests/record.py b/python/pyfory/tests/record.py
index 00546169f..ec929f298 100644
--- a/python/pyfory/tests/record.py
+++ b/python/pyfory/tests/record.py
@@ -20,9 +20,10 @@ from dataclasses import dataclass
import pyfory as fory
from typing import List, Dict
-from pyfory.utils import lazy_import
-
-pa = lazy_import("pyarrow")
+try:
+ import pyarrow as pa
+except ImportError:
+ pa = None
class Foo:
diff --git a/python/pyfory/tests/test_buffer.py
b/python/pyfory/tests/test_buffer.py
index f13590a2d..47a81388a 100644
--- a/python/pyfory/tests/test_buffer.py
+++ b/python/pyfory/tests/test_buffer.py
@@ -21,11 +21,9 @@ import pytest
import pyfory
from pyfory.serialization import ENABLE_FORY_CYTHON_SERIALIZATION, Buffer
-from pyfory.tests.core import require_pyarrow
+from pyfory.tests.core import pa, require_pyarrow
from pyfory.tests.test_stream import OneByteStream
-from pyfory.utils import clear_bit, get_bit, lazy_import, set_bit, set_bit_to
-
-pa = lazy_import("pyarrow")
+from pyfory.utils import clear_bit, get_bit, set_bit, set_bit_to
class RecvIntoOnlyStream:
diff --git a/python/pyfory/tests/test_cross_language.py
b/python/pyfory/tests/test_cross_language.py
index e640b92bf..68a335776 100644
--- a/python/pyfory/tests/test_cross_language.py
+++ b/python/pyfory/tests/test_cross_language.py
@@ -30,14 +30,11 @@ try:
except ImportError:
pytest = None
from dataclasses import dataclass
-from pyfory.utils import lazy_import
from typing import List, Dict, Any
import numpy as np
-pa = lazy_import("pyarrow")
-
def debug_print(*params):
"""print params if debug is needed."""
@@ -213,6 +210,8 @@ def test_serialization_with_schema(schema_file_path,
data_file_path):
@cross_language_test
def test_record_batch_basic(data_file_path):
+ import pyarrow as pa
+
with open(data_file_path, "rb") as f:
record_batch_bytes = f.read()
buf = pa.py_buffer(record_batch_bytes)
@@ -225,6 +224,8 @@ def test_record_batch_basic(data_file_path):
@cross_language_test
def test_record_batch(data_file_path):
+ import pyarrow as pa
+
with open(data_file_path, "rb") as f:
record_batch_bytes = f.read()
buf = pa.py_buffer(record_batch_bytes)
@@ -250,6 +251,8 @@ def test_record_batch(data_file_path):
@cross_language_test
def test_write_multi_record_batch(schema_file_path, data_file_path):
+ import pyarrow as pa
+
with open(schema_file_path, "rb") as f:
schema_bytes = f.read()
schema = pa.ipc.read_schema(pa.py_buffer(schema_bytes))
diff --git a/python/pyfory/tests/test_serializer.py
b/python/pyfory/tests/test_serializer.py
index 881561098..bb1d55cf0 100644
--- a/python/pyfory/tests/test_serializer.py
+++ b/python/pyfory/tests/test_serializer.py
@@ -46,9 +46,7 @@ from pyfory.serializer import (
Numpy1DArraySerializer,
)
from pyfory.types import TypeId
-from pyfory.utils import lazy_import
-
-pa = lazy_import("pyarrow")
+from pyfory.tests.core import pa
def test_compatible_mode_overrides():
diff --git a/python/pyfory/utils.py b/python/pyfory/utils.py
index c70082cef..2a262d7b7 100644
--- a/python/pyfory/utils.py
+++ b/python/pyfory/utils.py
@@ -15,69 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-import importlib
-import inspect
import sys
-from typing import Dict, Callable
-
-
-# This method is derived from
https://github.com/mars-project/mars/blob/c36c53fa22e10ef9477d9c454401a2f281375f31/mars/utils.py.
-def lazy_import(
- name: str,
- package: str = None,
- globals_: Dict = None, # pylint: disable=redefined-builtin
- locals_: Dict = None, # pylint: disable=redefined-builtin
- rename: str = None,
- placeholder: bool = False,
-):
- rename = rename or name
- prefix_name = name.split(".", 1)[0]
- globals_ = globals_ or inspect.currentframe().f_back.f_globals
-
- class LazyModule(object):
- def __init__(self):
- self._on_loads = []
-
- def __getattr__(self, item):
- if item.startswith("_pytest") or item in ("__bases__", "__test__"):
- raise AttributeError(item)
-
- real_mod = importlib.import_module(name, package=package)
- if rename in globals_:
- globals_[rename] = real_mod
- elif locals_ is not None:
- locals_[rename] = real_mod
- ret = getattr(real_mod, item)
- for on_load_func in self._on_loads:
- on_load_func()
- # make sure on_load hooks only executed once
- self._on_loads = []
- return ret
-
- def add_load_handler(self, func: Callable):
- self._on_loads.append(func)
- return func
-
- if importlib.util.find_spec(prefix_name) is not None:
- return LazyModule()
- elif placeholder:
- return ModulePlaceholder(prefix_name)
- else:
- return None
-
-
-class ModulePlaceholder:
- def __init__(self, mod_name: str):
- self._mod_name = mod_name
-
- def _raises(self):
- raise AttributeError(f"{self._mod_name} is required but not
installed.")
-
- def __getattr__(self, key):
- self._raises()
-
- def __call__(self, *_args, **_kwargs):
- self._raises()
is_little_endian = sys.byteorder == "little"
@@ -112,6 +50,5 @@ __all__ = [
"set_bit",
"clear_bit",
"set_bit_to",
- "lazy_import",
"is_little_endian",
]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
