Use pa.array constructor for extension types

apache · Oct 22, 2024 · e8baa24 · e8baa24
1 parent 20fe633
commit e8baa24
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/python/pyarrow/tests/parquet/test_data_types.py b/python/pyarrow/tests/parquet/test_data_types.py
@@ -516,10 +516,8 @@ def test_large_binary_overflow():
     pa.string(), pa.large_string()))
 def test_json_extension_type(storage_type):
     data = ['{"a": 1}', '{"b": 2}', None]
-    storage = pa.array(data, type=storage_type)
-    json_type = pa.json_(storage_type)
+    arr = pa.array(data, type=pa.json_(storage_type))
 
-    arr = pa.ExtensionArray.from_storage(json_type, storage)
     table = pa.table([arr], names=["ext"])
 
     _simple_table_roundtrip(table)
diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py
@@ -1932,8 +1932,9 @@ def test_bool8_scalar():
     pa.string(), pa.large_string(), pa.string_view()))
 def test_json(storage_type, pickle_module):
     data = ['{"a": 1}', '{"b": 2}', None]
-    storage = pa.array(data, type=storage_type)
     json_type = pa.json_(storage_type)
+    storage = pa.array(data, type=storage_type)
+    array = pa.array(data, type=json_type)
     json_arr_class = json_type.__arrow_ext_class__()
 
     assert pa.json_() == pa.json_(pa.utf8())
@@ -1944,7 +1945,6 @@ def test_json(storage_type, pickle_module):
     assert json_type == pa.json_(storage_type)
     assert json_type != storage_type
 
-    array = pa.ExtensionArray.from_storage(json_type, storage)
     assert isinstance(array, pa.JsonArray)
 
     assert array.to_pylist() == data