My personal benchmark shows Rust interface is much slower than Python interface

[ScottLinnn]

I'm writing 40 million 2-column rows to a dataset, they yield about 1.1GB size in total. I found that Rust is much slower than Python(2x time usage), which is a bit hard to understand. The code looks like below:

Python:

ROW_NUM = int(os.environ.get("BENCH_NUM_ROWS"))

def generate_dummy_data():
    # Two columns, one int - id, one varchar - name

    id_list = [i for i in range(ROW_NUM)]
    name_list = [f"test{i}" for i in range(ROW_NUM)]

    id_array = pa.array(id_list)
    name_array = pa.array(name_list)

    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
        ]
    )

    record_batch = pa.RecordBatch.from_arrays([id_array, name_array], schema=schema)

    return record_batch

def test_write():
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
        ]
    )

    home_dir = os.environ.get("HOME")
    base = home_dir + "/lance/file_jni_benchmark/python/"
    dataset_url = "dataset_python.lance"

    record_batch = generate_dummy_data()
    start_time = round(time.time() * 1000)
    lance.write_dataset(record_batch, base + dataset_url, schema)
    end_time = round(time.time() * 1000)

    elapsed_time = end_time - start_time
    print("Write finished - Python native write_dataset interface")
    print(f"Time used for write: {elapsed_time} miliseconds")

Rust code looks like this:

fn generate_dummy_data() -> RecordBatch {
    let row_num: usize = env::var("BENCH_NUM_ROWS").unwrap().parse().unwrap();
    let id_array = (0..row_num).map(|i| Some(i as i64)).collect::<Vec<_>>();
    let id_array = Int64Array::from(id_array);
    let name_array = StringArray::from(
        (0..row_num)
            .map(|i| Some(format!("test{}", i)))
            .collect::<Vec<_>>(),
    );

    let schema = Schema::new(vec![
        arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
        arrow::datatypes::Field::new("name", arrow::datatypes::DataType::Utf8, false),
    ]);

    RecordBatch::try_new(
        Arc::new(schema.clone()),
        vec![Arc::new(id_array), Arc::new(name_array)],
    )
    .unwrap()
}

fn test_write() {
    let base = format!(
        "{}/lance/file_jni_benchmark/rust/",
        env::var("HOME").unwrap()
    );
    let dataset_url = "dataset_rust.lance";

    // If this dir exists, remove it
    if fs::metadata(format!("{}{}", base, dataset_url)).is_ok() {
        fs::remove_dir_all(format!("{}{}", base, dataset_url)).unwrap();
    }

    let arrow_schema = Schema::new(vec![
        arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
        arrow::datatypes::Field::new("name", arrow::datatypes::DataType::Utf8, false),
    ]);

    let schema = Arc::new(arrow_schema.clone());

    let batches = vec![generate_dummy_data()];
    // let batches = vec![RecordBatch::new_empty(schema.clone())];
    let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);

    let write_params = WriteParams::default();
    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let start_time = SystemTime::now();
        Dataset::write(reader, &(base.clone() + dataset_url), Some(write_params))
            .await
            .unwrap();
        let elapsed_time = start_time.elapsed().unwrap().as_millis();

        println!("Write finished - Rust native write_dataset interface");
        println!("Time used for write: {} milliseconds", elapsed_time);
    });
}

Rust uses 20917ms and Python uses 9510ms. Are there anything I'm doing wrong with using the interfaces? Thanks!

[wjones127]

Hi @ScottLinnn. It might be informative to look at our compiler config for the Python binaries:

lance/python/.cargo/config.toml

Line 25 in fb43192

rustflags = ["-C", "target-cpu=haswell", "-C", "target-feature=+avx2,+fma,+f16c"]

lance/python/.cargo/config.toml

Line 2 in fb43192

lto = "thin"

We enable certain flags to get SIMD support built in. We also enable "thin" LTO. Both those can have significant impact.

Also I'm not sure how impactful this would be for write, but in tokio, you want to make sure you've enabled the multi-threaded runtime.

lance/python/Cargo.toml

Line 45 in fb43192

tokio = { version = "1.23", features = ["rt-multi-thread"] }

[ScottLinnn]

Thanks!