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.

Sign up for GitHub

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

Jump to bottom

Closes #2855: Read batches for Parquet strings #2856

Merged
merged 10 commits into from
Nov 20, 2023
Merged

Closes #2855: Read batches for Parquet strings #2856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
3cbe045
Convert string bytes calculation
bmcdonald3 Nov 15, 2023
e06cb6c
Working but with bad names
bmcdonald3 Nov 16, 2023
a5e4287
Remove writeln
bmcdonald3 Nov 16, 2023
24a10a0
Clean up
bmcdonald3 Nov 16, 2023
712abb7
Split string and list checks
bmcdonald3 Nov 16, 2023
f588501
Clean up spacing
bmcdonald3 Nov 16, 2023
60478d3
Fix null string reads
bmcdonald3 Nov 16, 2023
6de57f6
Fix null string reading
bmcdonald3 Nov 16, 2023
f28c2ce
Fix multi column write issue with null values
bmcdonald3 Nov 17, 2023
9f923ca
Remove duplicate value creation pointed out by Pierce
bmcdonald3 Nov 20, 2023
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
77 changes: 53 additions & 24 deletions src/ArrowFunctions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ int cpp_getListType(const char* filename, const char* colname, char** errMsg) {
}
}

int64_t cpp_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, char** errMsg) {
int64_t cpp_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg) {
try {
int64_t ty = cpp_getType(filename, colname, errMsg);
int64_t dty; // used to store the type of data so we can handle lists
Expand Down Expand Up @@ -275,25 +275,46 @@ int64_t cpp_getStringColumnNumBytes(const char* filename, const char* colname, v
}
column_reader = row_group_reader->Column(idx);

int16_t definition_level;
parquet::ByteArrayReader* ba_reader =
static_cast<parquet::ByteArrayReader*>(column_reader.get());

int64_t numRead = 0;
while (ba_reader->HasNext() && numRead < numElems) {
parquet::ByteArray value;
(void)ba_reader->ReadBatch(1, &definition_level, nullptr, &value, &values_read);
if ((ty == ARROWLIST && definition_level == 3) || ty == ARROWSTRING) {
if(values_read > 0) {
offsets[i] = value.len + 1;
byteSize += value.len + 1;
numRead += values_read;
} else {
offsets[i] = 1;
byteSize+=1;
numRead+=1;
if((numElems - numRead) < batchSize)
batchSize = numElems-numRead;
std::vector<parquet::ByteArray> string_values(batchSize);
std::vector<int16_t> definition_level(batchSize);
(void)ba_reader->ReadBatch(batchSize, definition_level.data(), nullptr, string_values.data(), &values_read);
numRead += values_read;
if (ty == ARROWSTRING) {
auto numCols = file_metadata -> num_columns();
int string_index = 0;
for(int idx = 0; idx < definition_level.size(); idx++) {
stress-tess marked this conversation as resolved.
Show resolved Hide resolved
auto lvl = definition_level[idx];
if(lvl != 0 || numCols > 1) {
auto value = string_values[string_index];
offsets[i] = value.len + 1;
byteSize += value.len + 1;
string_index++;
} else {
offsets[i] = 1;
byteSize+=1;
}
i++;
}
} else if (ty == ARROWLIST) {
for(int string_index = 0; string_index < values_read; string_index++) {
auto level = definition_level[string_index];
auto value = string_values[string_index];
if(value.len != 0) {
offsets[i] = value.len + 1;
byteSize += value.len + 1;
} else {
offsets[i] = 1;
byteSize+=1;
}
i++;
}
i++;
}
}
}
Expand Down Expand Up @@ -795,22 +816,30 @@ int cpp_readColumnByName(const char* filename, void* chpl_arr, const char* colna
i+=values_read;
}
} else if(ty == ARROWSTRING) {
auto numCols = file_metadata -> num_columns();
int16_t definition_level; // nullable type and only reading single records in batch
auto chpl_ptr = (unsigned char*)chpl_arr;
parquet::ByteArrayReader* reader =
static_cast<parquet::ByteArrayReader*>(column_reader.get());

while (reader->HasNext()) {
parquet::ByteArray value;
(void)reader->ReadBatch(1, &definition_level, nullptr, &value, &values_read);
// if values_read is 0, that means that it was a null value
if(values_read > 0) {
for(int j = 0; j < value.len; j++) {
chpl_ptr[i] = value.ptr[j];
i++;
std::vector<parquet::ByteArray> string_values(batchSize);
std::vector<int16_t> definition_level(batchSize);
(void)reader->ReadBatch(batchSize, definition_level.data(), nullptr, string_values.data(), &values_read);

int string_index = 0;
for (int idx = 0; idx < definition_level.size(); idx++) {
auto lvl = definition_level[idx];
if(lvl > 0 || numCols > 1) {
auto value = string_values[string_index];
for(int j = 0; j < value.len; j++) {
chpl_ptr[i] = value.ptr[j];
i++;
}
string_index++;
}
i++; // skip one space so the strings are null terminated with a 0
}
i++; // skip one space so the strings are null terminated with a 0
}
} else if(ty == ARROWFLOAT) {
auto chpl_ptr = (double*)chpl_arr;
Expand Down Expand Up @@ -2063,8 +2092,8 @@ extern "C" {
errMsg);
}

int64_t c_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, char** errMsg) {
return cpp_getStringColumnNumBytes(filename, colname, chpl_offsets, numElems, startIdx, errMsg);
int64_t c_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg) {
return cpp_getStringColumnNumBytes(filename, colname, chpl_offsets, numElems, startIdx, batchSize, errMsg);
}

int64_t c_getListColumnSize(const char* filename, const char* colname, void* chpl_seg_sizes, int64_t numElems, int64_t startIdx, char** errMsg) {
Expand Down
8 changes: 4 additions & 4 deletions src/ArrowFunctions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,10 @@ extern "C" {
const char* colname, int64_t numElems,
int64_t startIdx, int64_t batchSize, char** errMsg);

int64_t cpp_getStringColumnNumBytes(const char* filename, const char* colname,
void* chpl_offsets, int64_t numElems, int64_t startIdx, char** errMsg);
int64_t c_getStringColumnNumBytes(const char* filename, const char* colname,
void* chpl_offsets, int64_t numElems, int64_t startIdx, char** errMsg);
int64_t cpp_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets,
int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg);
int64_t c_getStringColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets,
int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg);

int64_t c_getListColumnSize(const char* filename, const char* colname,
void* chpl_seg_sizes, int64_t numElems, int64_t startIdx, char** errMsg);
Expand Down
4 changes: 2 additions & 2 deletions src/ParquetMsg.chpl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,13 +296,13 @@ module ParquetMsg {
}

proc getStrColSize(filename: string, dsetname: string, ref offsets: [] int) throws {
extern proc c_getStringColumnNumBytes(filename, colname, offsets, numElems, startIdx, errMsg): int;
extern proc c_getStringColumnNumBytes(filename, colname, offsets, numElems, startIdx, batchSize, errMsg): int;
var pqErr = new parquetErrorMsg();

var byteSize = c_getStringColumnNumBytes(filename.localize().c_str(),
dsetname.localize().c_str(),
c_ptrTo(offsets),
offsets.size, 0,
offsets.size, 0, 256,
c_ptrTo(pqErr.errMsg));

if byteSize == ARROWERROR then
Expand Down