Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closes #2880: Add parallel writing when writing pdarrays to Parquet #2881

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 29 additions & 4 deletions src/ParquetMsg.chpl
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ module ParquetMsg {
// Undocumented for now, just for internal experiments
private config const batchSize = getEnvInt("ARKOUDA_SERVER_PARQUET_BATCH_SIZE", 8192);

private config const parallelWriteThreshold = 512*1024*1024 / numBytes(int);

extern var ARROWINT64: c_int;
extern var ARROWINT32: c_int;
extern var ARROWUINT64: c_int;
Expand Down Expand Up @@ -417,6 +419,7 @@ module ParquetMsg {
dtype, compression,
errMsg): int;
var dtypeRep = toCDtype(dtype);
var doParallel = if A.size > parallelWriteThreshold then true else false;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

okay this is pedantic, but this could just be var doParallel = A.size > parallelWriteThreshold; right?

var prefix: string;
var extension: string;

Expand Down Expand Up @@ -454,10 +457,32 @@ module ParquetMsg {
valPtr = c_ptrTo(locArr);
}
if mode == TRUNCATE || !filesExist {
if c_writeColumnToParquet(myFilename.localize().c_str(), valPtr, 0,
dsetname.localize().c_str(), locDom.size, rowGroupSize,
dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR {
pqErr.parquetError(getLineNumber(), getRoutineName(), getModuleName());
if !doParallel {
if c_writeColumnToParquet(myFilename.localize().c_str(), valPtr, 0,
dsetname.localize().c_str(), locDom.size, rowGroupSize,
dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR {
pqErr.parquetError(getLineNumber(), getRoutineName(), getModuleName());
}
} else {
var fileSizes: [0..#loc.maxTaskPar] int = locDom.size/loc.maxTaskPar;
// First file has the extra elements if it isn't evenly divisible by maxTaskPar
fileSizes[0] += locDom.size - ((locDom.size/loc.maxTaskPar)*loc.maxTaskPar);
Comment on lines +467 to +469
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i had to convince myself that the integer cast of the float divide would always round down. This seems to check out and the adjustment gave what i expected on a small example!

var fileSizes: [0..#6] int = 10/6;  // 1.6666 verify this rounds down
writeln(fileSizes);
var leftOver = 10 - ((10/6)*6); 
writeln(leftOver);
1 1 1 1 1 1
4

I will say in my small example this resulted in a pretty unbalanced distribution, but I think that in a real case that locDom.size would be large enough relative to loc.maxTaskPar that it would be pretty uniform... I'm just now realizing that's probably part of the motivation for having a parallelWriteThreshold lol

Copy link
Member

@stress-tess stress-tess Dec 13, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since leftOver is the remainder of locDom.size/loc.maxTaskPar, it should be guaranteed to be less than fileSizes.size. So if we wanted to distribute the remainder more evenly we could do something like

var fileSizes: [0..#6] int = 10/6;
writeln(fileSizes);
var leftOver = 10 - ((10/6)*6); 
writeln(leftOver);
fileSizes[0..#leftOver] += 1;
writeln(fileSizes);
1 1 1 1 1 1
4
2 2 2 2 1 1

But this is probably overengineering something that isn't actually a problem


var offsets = + scan fileSizes;

forall i in fileSizes.domain {
var suffix = '%04i'.format(idx): string;
var parSuffix = '%04i'.format(i): string;
const parFilename = filename + "_LOCALE" + suffix + "_CORE" + parSuffix + ".parquet";
var oi = if i == 0 then i else offsets[i-1];
Copy link
Member

@stress-tess stress-tess Dec 13, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

instead of doing this back by one, couldn't you do

var offsets = (+ scan fileSizes) - fileSizes;
forall (i, off, len) in zip(fileSizes.domain, offsets, fileSizes) {
...

I don't think this would have any performance difference, but this is more similar to how we calculate offsets in other places. I normally prefer looping variables over indexing when possible because it makes easier for me to tell at a glance what's local, but that doesn't apply here. So there's def no need to change

var coreArr = locArr[oi..#(fileSizes[i])];
var pqErr = new parquetErrorMsg();
if c_writeColumnToParquet(parFilename.localize().c_str(), c_ptrTo(coreArr), 0,
dsetname.localize().c_str(), coreArr.size, rowGroupSize,
dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR {
pqErr.parquetError(getLineNumber(), getRoutineName(), getModuleName());
}
}
}
} else {
if c_appendColumnToParquet(myFilename.localize().c_str(), valPtr,
Expand Down
Loading