Method run

benchmarks/src/imdb/convert.rs:52–111

(self)

impl ConvertOpt {
    pub async fn run(self) -> Result<()> {
        let input_path = self.input_path.to_str().unwrap();
        let output_path = self.output_path.to_str().unwrap();
        let config = SessionConfig::new().with_batch_size(self.batch_size);
        let ctx = SessionContext::new_with_config(config);

        for table in IMDB_TABLES {
            let start = Instant::now();
            let schema = get_imdb_table_schema(table);
            let input_path = format!("{input_path}/{table}.csv");
            let output_path = format!("{output_path}/{table}.parquet");
            let options = CsvReadOptions::new()
                .schema(&schema)
                .has_header(false)
                .delimiter(b',')
                .escape(b'\\')
                .file_extension(".csv");

            let mut csv = ctx.read_csv(&input_path, options).await?;

            // Select all apart from the padding column
            let selection = csv
                .schema()
                .iter()
                .take(schema.fields.len())
                .map(Expr::from)
                .map(SelectExpr::from)
                .collect::<Vec<_>>();

            csv = csv.select(selection)?;

            println!(
                "Converting '{}' to {} files in directory '{}'",
                &input_path, self.file_format, &output_path
            );
            match self.file_format.as_str() {
                "csv" => {
                    csv.write_csv(
                        output_path.as_str(),
                        DataFrameWriteOptions::new(),
                        None,
                    )
                    .await?;
                }
                "parquet" => {
                    csv.write_parquet(
                        output_path.as_str(),
                        DataFrameWriteOptions::new(),
                        None,
                    )
                    .await?;
                }
                other => {
                    return not_impl_err!("Invalid output format: {other}");
                }
            }
            println!("Conversion completed in {} ms", start.elapsed().as_millis());
        }