-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathtest_struct_wal_bug.py
More file actions
186 lines (145 loc) Β· 6.05 KB
/
test_struct_wal_bug.py
File metadata and controls
186 lines (145 loc) Β· 6.05 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
"""
Test to replicate the struct column WAL bug.
The issue: When using orso as a WAL (append functionality + .arrow() call),
struct columns are converted to binary/blob in the schema's arrow_field property,
but the actual data remains as struct/dict. This causes a type mismatch.
"""
import pyarrow as pa
from pyarrow.lib import ArrowException
from orso import DataFrame
from orso.schema import FlatColumn
from orso.schema import RelationSchema
from orso.types import OrsoTypes
def test_struct_wal_without_schema():
"""Test appending struct data without explicit schema."""
print("\n=== Test 1: WAL without explicit schema ===")
# Create a DataFrame with initial data to establish schema
data = [
{"id": 1, "details": {"name": "Alice", "age": 30}},
{"id": 2, "details": {"name": "Bob", "age": 25}},
]
df = DataFrame(dictionaries=data)
# Append more data (simulating WAL usage)
df.append({"id": 3, "details": {"name": "Charlie", "age": 35}})
print("DataFrame schema:", df.schema)
print("DataFrame columns:", df.column_names)
# Try to convert to Arrow - this is where the bug should manifest
try:
arrow_table = df.arrow()
except ArrowException as error:
print("β Arrow conversion failed with error:")
print(f" {error.__class__.__name__}: {error}")
return False
print("β Arrow conversion succeeded!")
print("Arrow schema:", arrow_table.schema)
print("Arrow table:\n", arrow_table)
return True
def test_struct_wal_with_schema():
"""Test appending struct data with explicit schema."""
print("\n=== Test 2: WAL with explicit schema ===")
# Create a schema with a struct column
schema = RelationSchema(
name="test_schema",
columns=[
FlatColumn(name="id", type=OrsoTypes.INTEGER),
FlatColumn(name="details", type=OrsoTypes.STRUCT),
]
)
print("Schema columns:")
for col in schema.columns:
arrow_field = col.arrow_field
print(f" - {col.name}: {col.type} -> Arrow {arrow_field.type}")
# Create DataFrame with schema
df = DataFrame(schema=schema)
data = [
{"id": 1, "details": {"name": "Alice", "age": 30}},
{"id": 2, "details": {"name": "Bob", "age": 25}},
{"id": 3, "details": {"name": "Charlie", "age": 35}},
]
for row in data:
df.append(row)
print(f"\nDataFrame has {len(df)} rows")
# Try to convert to Arrow - this is where the bug should manifest
try:
arrow_table = df.arrow()
except ArrowException as error:
print("β Arrow conversion failed with error:")
print(f" {error.__class__.__name__}: {error}")
return False
print("β Arrow conversion succeeded!")
print("Arrow schema:", arrow_table.schema)
print("Arrow table:\n", arrow_table)
return True
def test_struct_from_arrow_roundtrip():
"""Test creating DataFrame from Arrow table with struct, then converting back."""
print("\n=== Test 3: Arrow -> Orso -> Arrow roundtrip ===")
# Create an Arrow table with a struct column
struct_type = pa.struct([
pa.field('name', pa.string()),
pa.field('age', pa.int64())
])
arrow_schema = pa.schema([
pa.field('id', pa.int64()),
pa.field('details', struct_type)
])
data = [
[1, 2, 3],
[{'name': 'Alice', 'age': 30}, {'name': 'Bob', 'age': 25}, {'name': 'Charlie', 'age': 35}]
]
original_table = pa.table(data, schema=arrow_schema)
print(f"Original Arrow table schema: {original_table.schema}")
print(f"Original Arrow table:\n{original_table}")
# Convert to Orso DataFrame
try:
df = DataFrame.from_arrow(original_table)
print("\nβ Created DataFrame from Arrow")
print("DataFrame schema:", df.schema)
arrow_table = df.arrow()
except ArrowException as error:
print("β Conversion failed with error:")
print(f" {error.__class__.__name__}: {error}")
return False
print("β Converted back to Arrow!")
print("Converted Arrow schema:", arrow_table.schema)
print("Converted Arrow table:\n", arrow_table)
return True
def test_schema_arrow_field_for_struct():
"""Test what arrow_field property returns for struct columns."""
print("\n=== Test 4: Schema arrow_field property for STRUCT ===")
# Create a FlatColumn with STRUCT type
struct_col = FlatColumn(name="details", type=OrsoTypes.STRUCT)
fallback_type = struct_col.arrow_field.type
nested_struct = FlatColumn(
name="details",
type=OrsoTypes.STRUCT,
fields=[
FlatColumn(name="name", type=OrsoTypes.VARCHAR),
FlatColumn(name="age", type=OrsoTypes.INTEGER),
],
)
nested_type = nested_struct.arrow_field.type
print("Fallback arrow type:", fallback_type)
print("Nested arrow type:", nested_type)
return fallback_type == pa.binary() and nested_type == pa.struct(
[pa.field("name", pa.string()), pa.field("age", pa.int64())]
)
if __name__ == "__main__":
print("=" * 60)
print("Testing struct column WAL bug")
print("=" * 60)
test4_result = test_schema_arrow_field_for_struct()
test1_result = test_struct_wal_without_schema()
test2_result = test_struct_wal_with_schema()
test3_result = test_struct_from_arrow_roundtrip()
print("\n" + "=" * 60)
print("SUMMARY:")
print("=" * 60)
print(f"Test 4 (STRUCT arrow_field behavior): {'β PASS' if test4_result else 'β FAIL'}")
print(f"Test 1 (WAL without schema): {'β PASS' if test1_result else 'β FAIL'}")
print(f"Test 2 (WAL with schema): {'β PASS' if test2_result else 'β FAIL'}")
print(f"Test 3 (Arrow roundtrip): {'β PASS' if test3_result else 'β FAIL'}")
print("=" * 60)
if all([test1_result, test2_result, test3_result, test4_result]):
print("\nβ All tests passed - structs now round-trip as structs")
else:
print("\nβ Mixed results - further investigation needed")