|
|
| """Test_PyArrow.py - on-device validation for the pyarrow 25.0.1 wheel.
|
|
|
| Exercises arrays, tables, IPC, CSV, JSON, Feather, compute, and types.
|
| Exit code 0 = all tests passed.
|
|
|
| Generated by RIMI
|
| """
|
| import os
|
| import sys
|
| import tempfile
|
|
|
| RESULTS = []
|
|
|
|
|
| def test(name, fn):
|
| try:
|
| fn()
|
| RESULTS.append(("PASS", name))
|
| except NotImplementedError:
|
| RESULTS.append(("SKIP", name))
|
| except Exception as e:
|
| RESULTS.append(("FAIL", name, str(e)))
|
|
|
|
|
| def section(title):
|
| print("\n===== %s =====" % title)
|
|
|
|
|
| def check(cond, msg="assertion failed"):
|
| if not cond:
|
| raise AssertionError(msg)
|
|
|
|
|
| WORKDIR = None
|
|
|
|
|
| def workdir():
|
| global WORKDIR
|
| if WORKDIR is None:
|
| import os as _os
|
| candidates = [_os.environ.get("TMPDIR") or "", tempfile.gettempdir(),
|
| "/storage/emulated/0/Download", _os.getcwd()]
|
| for base in candidates:
|
| if not base:
|
| continue
|
| try:
|
| d = os.path.join(base, "test_pyarrow_tmp")
|
| os.makedirs(d, exist_ok=True)
|
| with open(os.path.join(d, "_probe"), "w") as fh:
|
| fh.write("ok")
|
| WORKDIR = d
|
| break
|
| except OSError:
|
| continue
|
| if WORKDIR is None:
|
| WORKDIR = "."
|
| return WORKDIR
|
|
|
|
|
|
|
|
|
|
|
| def test_import_pyarrow():
|
| import pyarrow as pa
|
| print(" pyarrow", pa.__version__)
|
| check(pa.__version__.startswith("25."), "unexpected version: %s" % pa.__version__)
|
|
|
|
|
| def test_import_numpy_dep():
|
| import numpy as np
|
| print(" numpy", np.__version__)
|
| check(hasattr(np, "ndarray"), "numpy not functional")
|
|
|
|
|
| def test_import_submodules():
|
| import pyarrow.compute as pc
|
| import pyarrow.csv
|
| import pyarrow.json
|
| import pyarrow.feather
|
| import pyarrow.fs
|
| import pyarrow.ipc
|
| check(hasattr(pc, "cast"), "compute module incomplete")
|
| print(" compute, csv, json, feather, fs, ipc -- all imported")
|
|
|
|
|
|
|
|
|
|
|
| def test_array_creation():
|
| import pyarrow as pa
|
| a = pa.array([1, 2, 3, 4, 5])
|
| check(a.type == pa.int64(), "type %r" % a.type)
|
| check(len(a) == 5, "len %d" % len(a))
|
| check(a.to_pylist() == [1, 2, 3, 4, 5], "values mismatch")
|
| print(" int64 array:", a)
|
|
|
|
|
| def test_array_float():
|
| import pyarrow as pa
|
| a = pa.array([1.0, 2.5, 3.7], type=pa.float64())
|
| check(a.type == pa.float64(), "type %r" % a.type)
|
| check(len(a) == 3, "len %d" % len(a))
|
| print(" float64 array:", a)
|
|
|
|
|
| def test_array_string():
|
| import pyarrow as pa
|
| a = pa.array(["hello", "world", "pyarrow"])
|
| check(a.type == pa.string(), "type %r" % a.type)
|
| check(a.to_pylist() == ["hello", "world", "pyarrow"])
|
| print(" string array:", a)
|
|
|
|
|
| def test_array_null():
|
| import pyarrow as pa
|
| a = pa.array([1, None, 3], type=pa.int64())
|
| check(a.null_count == 1, "null_count %d" % a.null_count)
|
| check(a.to_pylist() == [1, None, 3])
|
| print(" null handling:", a)
|
|
|
|
|
|
|
|
|
|
|
| def test_types():
|
| import pyarrow as pa
|
| t_int = pa.int64()
|
| t_float = pa.float64()
|
| t_str = pa.string()
|
| t_bool = pa.bool_()
|
| check(str(t_int) == "int64", "int64 str: %r" % str(t_int))
|
| check(str(t_float) in ("float64", "double"), "float64 str: %r" % str(t_float))
|
| check(str(t_str) == "string", "string str: %r" % str(t_str))
|
| check(str(t_bool) == "bool", "bool str: %r" % str(t_bool))
|
| check(isinstance(t_int, pa.DataType), "not DataType")
|
| print(" types: int64, float64, string, bool -- all recognized")
|
|
|
|
|
|
|
|
|
|
|
| def test_table_creation():
|
| import pyarrow as pa
|
| t = pa.table({
|
| "id": [1, 2, 3],
|
| "name": ["alice", "bob", "charlie"],
|
| "score": [95.5, 87.0, 92.3],
|
| })
|
| check(t.num_rows == 3, "rows %d" % t.num_rows)
|
| check(t.num_columns == 3, "cols %d" % t.num_columns)
|
| check(t.column_names == ["id", "name", "score"])
|
| check(t.schema.field("id").type == pa.int64())
|
| check(t.schema.field("name").type == pa.string())
|
| check(t.schema.field("score").type == pa.float64())
|
| print(" table: %d rows x %d cols" % (t.num_rows, t.num_columns))
|
|
|
|
|
| def test_table_ops():
|
| import pyarrow as pa
|
| t = pa.table({"x": [10, 20, 30], "y": [1.0, 2.0, 3.0]})
|
| check(t.column("x").to_pylist() == [10, 20, 30])
|
| check(t.to_pandas().shape == (3, 2))
|
| print(" table column access + to_pandas OK")
|
|
|
|
|
|
|
|
|
|
|
| def test_ipc_roundtrip():
|
| import pyarrow as pa
|
| import pyarrow.ipc as ipc
|
| t = pa.table({
|
| "a": [1, 2, 3, 4, 5],
|
| "b": ["x", "y", "z", "w", "v"],
|
| "c": [1.1, 2.2, 3.3, 4.4, 5.5],
|
| })
|
| path = os.path.join(workdir(), "test_ipc.arrow")
|
| sink = ipc.new_file(path, t.schema)
|
| sink.write_table(t)
|
| sink.close()
|
| reader = ipc.open_file(path)
|
| t2 = reader.read_all()
|
| check(t.equals(t2), "IPC roundtrip mismatch")
|
| os.unlink(path)
|
| print(" IPC file: write -> read -> equals")
|
|
|
|
|
| def test_ipc_stream():
|
| import pyarrow as pa
|
| import pyarrow.ipc as ipc
|
| t = pa.table({"val": [100, 200, 300]})
|
| path = os.path.join(workdir(), "test_ipc_stream.arrow")
|
| sink = ipc.new_stream(path, t.schema)
|
| sink.write_table(t)
|
| sink.close()
|
| reader = ipc.open_stream(path)
|
| t2 = reader.read_all()
|
| check(t.equals(t2), "IPC stream roundtrip mismatch")
|
| os.unlink(path)
|
| print(" IPC stream: write -> read -> equals")
|
|
|
|
|
|
|
|
|
|
|
| def test_csv_write_read():
|
| import pyarrow as pa
|
| import pyarrow.csv as pcsv
|
| t = pa.table({
|
| "id": [1, 2, 3],
|
| "name": ["alice", "bob", "charlie"],
|
| "val": [10.5, 20.3, 30.1],
|
| })
|
| path = os.path.join(workdir(), "test_csv.csv")
|
| pcsv.write_csv(t, path)
|
| t2 = pcsv.read_csv(path)
|
| check(t2.num_rows == 3, "rows %d" % t2.num_rows)
|
| check(t2.num_columns == 3, "cols %d" % t2.num_columns)
|
| os.unlink(path)
|
| print(" CSV write -> read: %d rows x %d cols" % (t2.num_rows, t2.num_columns))
|
|
|
|
|
| def test_csv_options():
|
| import pyarrow as pa
|
| import pyarrow.csv as pcsv
|
| path = os.path.join(workdir(), "test_csv_opts.csv")
|
| with open(path, "w") as fh:
|
| fh.write("1,2,3\n4,5,6\n")
|
| read_opts = pcsv.ReadOptions(column_names=["x", "y", "z"])
|
| convert_opts = pcsv.ConvertOptions(column_types={"x": pa.int64(), "y": pa.int64(), "z": pa.int64()})
|
| t = pcsv.read_csv(path, read_options=read_opts, convert_options=convert_opts)
|
| check(t.column("x").to_pylist() == [1, 4])
|
| check(t.schema.field("x").type == pa.int64())
|
| os.unlink(path)
|
| print(" CSV with custom options: column_names + column_types")
|
|
|
|
|
|
|
|
|
|
|
| def test_json_read():
|
| import pyarrow as pa
|
| import pyarrow.json as pjson
|
| path = os.path.join(workdir(), "test_json.json")
|
| with open(path, "w") as fh:
|
| fh.write('{"a": 1, "b": "hello"}\n')
|
| fh.write('{"a": 2, "b": "world"}\n')
|
| t = pjson.read_json(path)
|
| check(t.num_rows == 2, "rows %d" % t.num_rows)
|
| check("a" in t.column_names, "column 'a' missing")
|
| check("b" in t.column_names, "column 'b' missing")
|
| os.unlink(path)
|
| print(" JSON read: %d rows, columns=%r" % (t.num_rows, t.column_names))
|
|
|
|
|
|
|
|
|
|
|
| def test_feather_roundtrip():
|
| import pyarrow as pa
|
| import pyarrow.feather as pf
|
| t = pa.table({
|
| "id": [1, 2, 3, 4, 5],
|
| "name": ["alice", "bob", "charlie", "diana", "eve"],
|
| "score": [95.5, 87.0, 92.3, 88.8, 99.1],
|
| })
|
| path = os.path.join(workdir(), "test_feather.feather")
|
| pf.write_feather(t, path)
|
| t2 = pf.read_table(path)
|
| check(t.equals(t2), "Feather roundtrip mismatch")
|
| check(t2.num_rows == 5, "rows %d" % t2.num_rows)
|
| os.unlink(path)
|
| print(" Feather: write -> read -> equals (%d rows)" % t2.num_rows)
|
|
|
|
|
|
|
|
|
|
|
| def test_compute_basic():
|
| import pyarrow as pa
|
| import pyarrow.compute as pc
|
| a = pa.array([1, 2, 3, 4, 5])
|
| result = pc.sum(a)
|
| check(result.as_py() == 15, "sum %r" % result)
|
| print(" compute.sum:", result.as_py())
|
|
|
|
|
| def test_compute_cast():
|
| import pyarrow as pa
|
| import pyarrow.compute as pc
|
| a = pa.array([1, 2, 3], type=pa.int64())
|
| b = pc.cast(a, pa.float64())
|
| check(b.type == pa.float64(), "type %r" % b.type)
|
| check(b.to_pylist() == [1.0, 2.0, 3.0])
|
| print(" compute.cast int64 -> float64:", b)
|
|
|
|
|
| def test_compute_filter():
|
| import pyarrow as pa
|
| import pyarrow.compute as pc
|
| a = pa.array([10, 20, 30, 40, 50])
|
| mask = pc.greater(a, 25)
|
| filtered = pc.filter(a, mask)
|
| check(filtered.to_pylist() == [30, 40, 50])
|
| print(" compute.filter > 25:", filtered)
|
|
|
|
|
| def test_compute_arithmetic():
|
| import pyarrow as pa
|
| import pyarrow.compute as pc
|
| a = pa.array([10, 20, 30])
|
| b = pa.array([1, 2, 3])
|
| add_result = pc.add(a, b)
|
| mul_result = pc.multiply(a, b)
|
| check(add_result.to_pylist() == [11, 22, 33])
|
| check(mul_result.to_pylist() == [10, 40, 90])
|
| print(" compute add/multiply:", add_result, mul_result)
|
|
|
|
|
|
|
|
|
|
|
| def test_fs_local():
|
| import pyarrow.fs as pfs
|
| local = pfs.LocalFileSystem()
|
| path = os.path.join(workdir(), "test_fs.txt")
|
| with open(path, "w") as fh:
|
| fh.write("filesystem test")
|
| meta = local.get_file_info(path)
|
| check(meta.type == pfs.FileType.File, "not a file")
|
| check(meta.size > 0, "size %d" % meta.size)
|
| os.unlink(path)
|
| print(" LocalFileSystem: get_file_info OK (size=%d)" % meta.size)
|
|
|
|
|
|
|
|
|
|
|
| def main():
|
| section("pyarrow 25.0.1 - import / basics")
|
| test("import pyarrow (25.x)", test_import_pyarrow)
|
| test("import numpy (dependency)", test_import_numpy_dep)
|
| test("import submodules (compute, csv, json, feather, fs, ipc)", test_import_submodules)
|
|
|
| section("arrays")
|
| test("array int64", test_array_creation)
|
| test("array float64", test_array_float)
|
| test("array string", test_array_string)
|
| test("array null handling", test_array_null)
|
|
|
| section("types")
|
| test("types (int64, string, float64, bool)", test_types)
|
|
|
| section("tables")
|
| test("table creation + schema", test_table_creation)
|
| test("table column access + to_pandas", test_table_ops)
|
|
|
| section("IPC")
|
| test("IPC file round-trip", test_ipc_roundtrip)
|
| test("IPC stream round-trip", test_ipc_stream)
|
|
|
| section("CSV")
|
| test("CSV write / read", test_csv_write_read)
|
| test("CSV custom options", test_csv_options)
|
|
|
| section("JSON")
|
| test("JSON read", test_json_read)
|
|
|
| section("Feather")
|
| test("Feather round-trip", test_feather_roundtrip)
|
|
|
| section("compute")
|
| test("compute.sum", test_compute_basic)
|
| test("compute.cast", test_compute_cast)
|
| test("compute.filter", test_compute_filter)
|
| test("compute arithmetic", test_compute_arithmetic)
|
|
|
| section("filesystem")
|
| test("LocalFileSystem get_file_info", test_fs_local)
|
|
|
| section("RESULT")
|
| n_ok = n_fail = n_skip = 0
|
| for r in RESULTS:
|
| status = r[0]
|
| if status == "PASS":
|
| n_ok += 1
|
| print(" OK %s" % r[1])
|
| elif status == "SKIP":
|
| n_skip += 1
|
| print(" SKIP %s" % r[1])
|
| else:
|
| n_fail += 1
|
| print(" FAIL %s: %s" % (r[1], r[2]))
|
| print("RESULT: %d ok, %d failed, %d skipped" % (n_ok, n_fail, n_skip))
|
| sys.exit(1 if n_fail else 0)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
