import os
import json
import pandas as pd
import numpy as np
import csv
from datasets import load_dataset
from openpyxl import load_workbook
csv.field_size_limit(10 * 1024 * 1024)

def stream_read_any(path, xlsx_sheet_name=None):
    if path.endswith('.jsonl'):
        with open(path, 'r') as f:
            for line in f:
                yield json.loads(line.strip())
    elif path.endswith('.parquet'):
        ds = load_dataset("parquet", data_files=path, split="train", streaming=True)
        for row in ds:
            yield row
    elif path.endswith('.csv'):
        with open(path, "r", encoding='utf-8-sig') as f: 
            reader = csv.DictReader(f) 
            for row in reader:
                yield row
    elif path.endswith(".xlsx"):
        wb = load_workbook(filename=path, read_only=True)
        ws = wb[xlsx_sheet_name] if xlsx_sheet_name else wb.active

        # 获取表头
        rows_iter = ws.iter_rows(values_only=True)
        headers = next(rows_iter)
        
        for row in rows_iter:
            # 每行构造成字典
            yield dict(zip(headers, row))
    else:
        raise ValueError("error type:{}".format(path))

if __name__ == "__main__":
    path = "test.parquet"

    new_datas = []
    iterator = stream_read_any(path)
    for data in iterator:
        new_datas.append(data)
