Repository: risingwavelabs/risingwave
File: e2e_test/s3/file_source.py (Python)

Code before:
  s3.credentials.access = 'hummockadmin',
  s3.credentials.secret = 'hummockadmin',
  s3.endpoint_url = 'http://hummock001.127.0.0.1:9301',
  refresh.interval.sec = 1
) FORMAT PLAIN ENCODE PARQUET;
"""
    )

    # Wait for file to be ingested.
    MAX_RETRIES = 40
    for retry_no in range(MAX_RETRIES):
        cur.execute(f"select count(*) from {table_name}")
        result = cur.fetchone()
        if result[0] == 2:
            break
        print(
            f"[retry {retry_no}] Now got {result[0]} rows in table, 2 expected, wait 5s"
        )
        sleep(5)

    # Verify we can read subset fields from struct.
    # RisingWave struct field access uses the Postgres-compatible syntax: (s).foo
    cur.execute(f"select id, (s).foo, (s).bar from {table_name} order by id")
    rows = cur.fetchall()
    assert rows == [(1, 10, "a"), (2, 20, "b")], f"unexpected rows: {rows}"
    print("Test parquet struct subset read pass")

    # Cleanup RW & MinIO & local.
    cur.execute(f"drop table {table_name};")
    cur.close()
    conn.close()

    minio_client.remove_object(bucket_name="hummock001", object_name=s3_key)
    try:
        os.remove(local_path)
    except OSError:
        pass

def gen_data(file_num, item_num_per_file):
    assert item_num_per_file % 2 == 0, \
        f'item_num_per_file should be even to ensure sum(mark) == 0: {item_num_per_file}'
    return [
        [{
            'id': file_id * item_num_per_file + item_id,
            'name': f'{file_id}_{item_id}',
            'sex': item_id % 2,
            'mark': (-1) ** (item_id % 2),
        } for item_id in range(item_num_per_file)]
        for file_id in range(file_num)
    ]


Review comment (suggestion):
can you also cover the case when the desired field not exist in the struct?
like in rw schema `s STRUCT<foo INT, bar VARCHAR>`, but in parquet only have `s STRUCT<foo INT>`

Code after:
  s3.credentials.access = 'hummockadmin',
  s3.credentials.secret = 'hummockadmin',
  s3.endpoint_url = 'http://hummock001.127.0.0.1:9301',
  refresh.interval.sec = 1
) FORMAT PLAIN ENCODE PARQUET;
"""
    )

    # Wait for file to be ingested.
    MAX_RETRIES = 40
    for retry_no in range(MAX_RETRIES):
        cur.execute(f"select count(*) from {table_name}")
        result = cur.fetchone()
        if result[0] == 2:
            break
        print(
            f"[retry {retry_no}] Now got {result[0]} rows in table, 2 expected, wait 5s"
        )
        sleep(5)

    # Verify we can read subset fields from struct.
    # RisingWave struct field access uses the Postgres-compatible syntax: (s).foo
    cur.execute(f"select id, (s).foo, (s).bar from {table_name} order by id")
    rows = cur.fetchall()
    assert rows == [(1, 10, "a"), (2, 20, "b")], f"unexpected rows: {rows}"
    print("Test parquet struct subset read pass")

    # Cleanup RW & MinIO & local.
    cur.execute(f"drop table {table_name};")
    cur.close()
    conn.close()

    minio_client.remove_object(bucket_name="hummock001", object_name=s3_key)
    try:
        os.remove(local_path)
    except OSError:
        pass


def _test_parquet_struct_missing_field_in_file(minio_client: Minio):
    """Generate a Parquet file whose struct is missing some fields declared in RW schema.

    Upstream Parquet schema:
      - id: int32
      - s: struct<foo:int32>

    RisingWave table schema (superset):
      - id: int
      - s: struct<foo:int, bar:varchar>

    Expectation: