Function test_ndjson_partitioned

datafusion/datasource-json/src/source.rs:877–932

()

    #[tokio::test]
    async fn test_ndjson_partitioned() -> Result<()> {
        // Build an NDJSON file with a known number of rows.
        let num_rows: usize = 20;
        let mut ndjson = String::new();
        for i in 0..num_rows {
            ndjson.push_str(&format!("{{\"id\": {i}, \"name\": \"user{i}\"}}\n"));
        }
        let ndjson_bytes = Bytes::from(ndjson);
        let file_size = ndjson_bytes.len() as u64;

        for &cs in CHUNK_SIZES {
            let (store, path) = make_chunked_store(&ndjson_bytes, cs).await;

            for num_partitions in get_partition_splits() {
                let batches = collect_partitioned_batches(
                    Arc::clone(&store),
                    &path,
                    file_size,
                    num_partitions,
                )
                .await?;

                let total: usize = batches.iter().map(|b| b.num_rows()).sum();
                assert_eq!(
                    total, num_rows,
                    "Expected {num_rows} rows with {num_partitions} partitions"
                );

                let result = concat_and_sort_by_id(&batches)?;
                let ids = result
                    .column(0)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap();
                let names = result
                    .column(1)
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .unwrap();
                for i in 0..num_rows {
                    assert_eq!(
                        ids.value(i),
                        i as i64,
                        "id mismatch at row {i} with {num_partitions} partitions"
                    );
                    assert_eq!(
                        names.value(i),
                        format!("user{i}"),
                        "name mismatch at row {i} with {num_partitions} partitions"
                    );
                }
            }
        }

        Ok(())
    }

    #[tokio::test]