()
| 36 | |
| 37 | @require_pyarrow |
| 38 | def test_vectorized(): |
| 39 | field_names = ["f" + str(i) for i in range(1, 6)] |
| 40 | cls = fory.record_class_factory("TEST_VECTORIZED", field_names) |
| 41 | # Create Fory schema for the encoder |
| 42 | fory_schema = schema( |
| 43 | [ |
| 44 | field("f1", int64()), |
| 45 | field("f2", int32()), |
| 46 | field("f3", int16()), |
| 47 | field("f4", int8()), |
| 48 | field("f5", utf8()), |
| 49 | ] |
| 50 | ) |
| 51 | # Convert to Arrow schema for ArrowWriter |
| 52 | arrow_schema = to_arrow_schema(fory_schema) |
| 53 | # Add metadata for class resolution |
| 54 | arrow_schema = arrow_schema.with_metadata({"cls": fory.get_qualified_classname(cls)}) |
| 55 | |
| 56 | writer = fory.format.ArrowWriter(arrow_schema) |
| 57 | encoder = fory.create_row_encoder(fory_schema) |
| 58 | num_rows = 10 |
| 59 | data = [[] for _ in range(len(field_names))] |
| 60 | for i in range(num_rows): |
| 61 | obj = cls(f1=2**63 - 1, f2=2**31 - 1, f3=2**15 - 1, f4=2**7 - 1, f5=f"str{i}") |
| 62 | fields_data = list(obj) |
| 63 | for j in range(len(fields_data)): |
| 64 | data[j].append(fields_data[j]) |
| 65 | row = encoder.to_row(obj) |
| 66 | writer.write(row) |
| 67 | record_batch = writer.finish() |
| 68 | writer.reset() |
| 69 | print(f"record_batch {record_batch}") |
| 70 | print(f"record_batch.num_rows {record_batch.num_rows}") |
| 71 | print(f"record_batch.num_columns {record_batch.num_columns}") |
| 72 | assert record_batch.num_rows == num_rows |
| 73 | assert record_batch.num_columns == 5 |
| 74 | |
| 75 | data = [pa.array(data[i], type=arrow_schema[i].type) for i in range(len(field_names))] |
| 76 | batch1 = pa.RecordBatch.from_arrays(data, field_names) |
| 77 | assert batch1 == record_batch |
| 78 | |
| 79 | batches = [record_batch] * 3 |
| 80 | table = pa.Table.from_batches(batches) |
| 81 | print(f"table {table}") |
| 82 | |
| 83 | |
| 84 | @require_pyarrow |
no test coverage detected