diff --git a/dep/checkArrow.chpl b/dep/checkArrow.chpl index f70b72736e5..cdf716d8af7 100644 --- a/dep/checkArrow.chpl +++ b/dep/checkArrow.chpl @@ -1,33 +1,6 @@ -use IO; -use CTypes; - -require "../src/parquet/WriteParquet.h"; -require "../src/WriteParquet.o"; -require "../src/parquet/ReadParquet.h"; -require "../src/ReadParquet.o"; -require "../src/parquet/UtilParquet.h"; -require "../src/UtilParquet.o"; - -proc getVersionInfo() { - extern proc c_getVersionInfo(): c_ptrConst(c_char); - extern proc strlen(str): c_int; - extern proc c_free_string(ptr); - var cVersionString = c_getVersionInfo(); - defer { - c_free_string(cVersionString: c_ptr(void)); - } - var ret: string; - try { - ret = string.createCopyingBuffer(cVersionString, - strlen(cVersionString)); - } catch e { - ret = "Error converting Arrow version message to Chapel string"; - } - return ret; -} +use Parquet; proc main() { - var ArrowVersion = getVersionInfo(); - writeln("Found Arrow version: ", ArrowVersion); + writeln("Found Arrow version: ", getVersionInfo()); return 0; } diff --git a/make/Arkouda.mk b/make/Arkouda.mk index e7fb292a418..295f2b6a8ba 100644 --- a/make/Arkouda.mk +++ b/make/Arkouda.mk @@ -46,15 +46,16 @@ endif SERVER_CONFIG_SCRIPT=$(ARKOUDA_SOURCE_DIR)/parseServerConfig.py # This is the main compilation statement section -$(ARKOUDA_MAIN_MODULE): check-deps register-commands $(ARROW_UTIL_O) $(ARROW_READ_O) $(ARROW_WRITE_O) $(ARKOUDA_SOURCES) $(ARKOUDA_MAKEFILES) +$(ARKOUDA_MAIN_MODULE): check-deps register-commands $(ARKOUDA_SOURCES) $(ARKOUDA_MAKEFILES) $(eval MOD_GEN_OUT=$(shell python3 $(SERVER_CONFIG_SCRIPT) $(ARKOUDA_CONFIG_FILE) $(ARKOUDA_SOURCE_DIR))) - $(CHPL) $(CHPL_DEBUG_FLAGS) $(CHPL_USER_FLAGS) $(PRINT_PASSES_FLAGS) $(REGEX_MAX_CAPTURES_FLAG) $(OPTIONAL_SERVER_FLAGS) $(CHPL_FLAGS_WITH_VERSION) $(CHPL_COMPAT_FLAGS) $(ARKOUDA_MAIN_SOURCE) $(ARKOUDA_COMPAT_MODULES) $(ARKOUDA_SERVER_USER_MODULES) $(MOD_GEN_OUT) $(ARKOUDA_RW_DEFAULT_FLAG) $(ARKOUDA_KEYPART_FLAG) $(ARKOUDA_REGISTRY_DIR)/Commands.chpl -I$(ARKOUDA_SOURCE_DIR)/parquet -o $@ + PARQUET_PKG_FLAGS="$$($(PARQUET_PACKAGE_FLAGS_CMD))" && \ + $(CHPL) $(CHPL_DEBUG_FLAGS) $(CHPL_USER_FLAGS) $(PRINT_PASSES_FLAGS) $(REGEX_MAX_CAPTURES_FLAG) $(OPTIONAL_SERVER_FLAGS) $(CHPL_FLAGS_WITH_VERSION) $(CHPL_COMPAT_FLAGS) $(ARKOUDA_MAIN_SOURCE) $(ARKOUDA_COMPAT_MODULES) $(ARKOUDA_SERVER_USER_MODULES) $(MOD_GEN_OUT) $(ARKOUDA_RW_DEFAULT_FLAG) $(ARKOUDA_KEYPART_FLAG) $(ARKOUDA_REGISTRY_DIR)/Commands.chpl $$PARQUET_PKG_FLAGS -o $@ CLEAN_TARGETS += arkouda-clean .PHONY: arkouda-clean arkouda-clean: - $(RM) $(ARKOUDA_MAIN_MODULE) $(ARKOUDA_MAIN_MODULE)_real $(ARROW_UTIL_O) $(ARROW_READ_O) $(ARROW_WRITE_O) + $(RM) $(ARKOUDA_MAIN_MODULE) $(ARKOUDA_MAIN_MODULE)_real .PHONY: tags tags: diff --git a/make/Prologue.mk b/make/Prologue.mk index 31b50cf625c..bf17e5a671d 100644 --- a/make/Prologue.mk +++ b/make/Prologue.mk @@ -5,10 +5,10 @@ include make/prologue/core.mk include make/prologue/chapel.mk include make/prologue/paths.mk include make/prologue/config.mk -include make/prologue/arrow_shims.mk # Dependency meta-targets and per-dependency rules include make/deps/common.mk +include make/deps/parquet.mk include make/deps/zmq.mk include make/deps/hdf5.mk include make/deps/arrow.mk diff --git a/make/deps/parquet.mk b/make/deps/parquet.mk new file mode 100644 index 00000000000..f0e86b3a31d --- /dev/null +++ b/make/deps/parquet.mk @@ -0,0 +1,12 @@ +# External Chapel `Parquet` Mason package integration. +# +# get_parquet_package.sh clones/builds the package and prints the chpl arguments +# (C++ prerequisite headers/objects, include paths, Arrow link flags) plus the +# module source path. We use this instead of `mason modules`, which does not +# work reliably with Chapel 2.4. Override PARQUET_INSTALL_DIR to relocate the +# clone; set ARKOUDA_PARQUET_SRC_DIR to reuse an existing checkout. +PARQUET_PACKAGE_SCRIPT := $(ARKOUDA_PROJECT_DIR)/scripts/get_parquet_package.sh +PARQUET_INSTALL_DIR ?= $(DEP_BUILD_DIR)/Parquet + +PARQUET_PACKAGE_FLAGS_CMD = env ARKOUDA_CHPL_HOME="$(ARKOUDA_CHPL_HOME)" \ + "$(PARQUET_PACKAGE_SCRIPT)" "$(PARQUET_INSTALL_DIR)" diff --git a/make/prologue/arrow_shims.mk b/make/prologue/arrow_shims.mk deleted file mode 100644 index 6090361e9c1..00000000000 --- a/make/prologue/arrow_shims.mk +++ /dev/null @@ -1,59 +0,0 @@ -ARROW_PIC ?= - -ARROW_READ_FILE_NAME += $(ARKOUDA_SOURCE_DIR)/parquet/ReadParquet -ARROW_READ_CPP += $(ARROW_READ_FILE_NAME).cpp -ARROW_READ_H += $(ARROW_READ_FILE_NAME).h -ARROW_READ_O += $(ARKOUDA_SOURCE_DIR)/ReadParquet.o - -ARROW_WRITE_FILE_NAME += $(ARKOUDA_SOURCE_DIR)/parquet/WriteParquet -ARROW_WRITE_CPP += $(ARROW_WRITE_FILE_NAME).cpp -ARROW_WRITE_H += $(ARROW_WRITE_FILE_NAME).h -ARROW_WRITE_O += $(ARKOUDA_SOURCE_DIR)/WriteParquet.o - -ARROW_UTIL_FILE_NAME += $(ARKOUDA_SOURCE_DIR)/parquet/UtilParquet -ARROW_UTIL_CPP += $(ARROW_UTIL_FILE_NAME).cpp -ARROW_UTIL_H += $(ARROW_UTIL_FILE_NAME).h -ARROW_UTIL_O += $(ARKOUDA_SOURCE_DIR)/UtilParquet.o - -CHPL_CXX = $(shell $(ARKOUDA_CHPL_HOME)/util/config/compileline --compile-c++ 2>/dev/null) -ifeq ($(CHPL_CXX),) -CHPL_CXX=$(CXX) -endif - -SANITIZER = $(shell $(ARKOUDA_CHPL_HOME)/util/chplenv/chpl_sanitizers.py --exe 2>/dev/null) -ifneq ($(SANITIZER),none) -ARROW_SANITIZE = -fsanitize=$(SANITIZER) -endif - -.PHONY: \ - compile-arrow-cpp \ - compile-arrow-write \ - compile-arrow-read \ - compile-arrow-util - -compile-arrow-cpp: - $(MAKE) compile-arrow-write - $(MAKE) compile-arrow-read - $(MAKE) compile-arrow-util - -compile-arrow-write: - $(CHPL_CXX) -O3 -std=c++20 $(ARROW_PIC) -c $(ARROW_WRITE_CPP) -o $(ARROW_WRITE_O) \ - $(INCLUDE_FLAGS) $(ARROW_SANITIZE) - -compile-arrow-read: - $(CHPL_CXX) -O3 -std=c++20 $(ARROW_PIC) -c $(ARROW_READ_CPP) -o $(ARROW_READ_O) \ - $(INCLUDE_FLAGS) $(ARROW_SANITIZE) - -compile-arrow-util: - $(CHPL_CXX) -O3 -std=c++20 $(ARROW_PIC) -c $(ARROW_UTIL_CPP) -o $(ARROW_UTIL_O) \ - $(INCLUDE_FLAGS) $(ARROW_SANITIZE) - -$(ARROW_UTIL_O): $(ARROW_UTIL_CPP) $(ARROW_UTIL_H) - $(MAKE) compile-arrow-util - -$(ARROW_READ_O): $(ARROW_READ_CPP) $(ARROW_READ_H) - $(MAKE) compile-arrow-read - -$(ARROW_WRITE_O): $(ARROW_WRITE_CPP) $(ARROW_WRITE_H) - $(MAKE) compile-arrow-write - diff --git a/make/prologue/checks.mk b/make/prologue/checks.mk index 596c190daae..d905bf55a4c 100644 --- a/make/prologue/checks.mk +++ b/make/prologue/checks.mk @@ -35,15 +35,14 @@ check-re2: $(RE2_CHECK) @rm -f $(DEP_INSTALL_DIR)/$@ $(DEP_INSTALL_DIR)/$@_real ARROW_CHECK = $(DEP_INSTALL_DIR)/checkArrow.chpl -check-arrow: $(ARROW_CHECK) $(ARROW_UTIL_O) $(ARROW_READ_O) $(ARROW_WRITE_O) +check-arrow: $(ARROW_CHECK) @echo "Checking for Arrow" - $(MAKE) compile-arrow-cpp --no-print-directory - @$(CHPL) $(CHPL_FLAGS) $(ARKOUDA_COMPAT_MODULES) $< \ - $(ARROW_M) -M $(ARKOUDA_SOURCE_DIR) \ - -I $(ARKOUDA_SOURCE_DIR)/parquet \ + @PARQUET_PKG_FLAGS="$$($(PARQUET_PACKAGE_FLAGS_CMD))" && \ + $(CHPL) $(CHPL_FLAGS) $(ARKOUDA_COMPAT_MODULES) $< \ + $$PARQUET_PKG_FLAGS \ -o $(DEP_INSTALL_DIR)/$@ && \ - ([ $$? -eq 0 ] && echo "Success compiling program") || \ - echo "\nERROR: Please ensure that dependencies have been installed correctly (see -> https://github.com/Bears-R-Us/arkouda/blob/main/pydoc/setup/BUILD.md)\n" + echo "Success compiling program" || \ + { echo "\nERROR: Please ensure that dependencies have been installed correctly (see -> https://github.com/Bears-R-Us/arkouda/blob/main/pydoc/setup/BUILD.md)\n"; exit 1; } $(DEP_INSTALL_DIR)/$@ -nl 1 @rm -f $(DEP_INSTALL_DIR)/$@ $(DEP_INSTALL_DIR)/$@_real diff --git a/make/prologue/paths.mk b/make/prologue/paths.mk index 741ff262cee..6a9ba027358 100644 --- a/make/prologue/paths.mk +++ b/make/prologue/paths.mk @@ -9,10 +9,20 @@ ifneq ("$(wildcard $(1)/lib64)","") CHPL_FLAGS += -L$(1)/lib64 --ldflags="-Wl,-rpath,$(1)/lib64" endif +ifneq ("$(wildcard $(1)/lib64/pkgconfig)","") +PKG_CONFIG_PATH := $(1)/lib64/pkgconfig$$(if $$(PKG_CONFIG_PATH),:$$(PKG_CONFIG_PATH)) +endif + # Add lib if present ifneq ("$(wildcard $(1)/lib)","") CHPL_FLAGS += -L$(1)/lib --ldflags="-Wl,-rpath,$(1)/lib" endif + +ifneq ("$(wildcard $(1)/lib/pkgconfig)","") +PKG_CONFIG_PATH := $(1)/lib/pkgconfig$$(if $$(PKG_CONFIG_PATH),:$$(PKG_CONFIG_PATH)) +endif + +export PKG_CONFIG_PATH endef # Usage: $(eval $(call add-path,/home/user/anaconda3/envs/arkouda)) # ^ no space after comma diff --git a/pydoc/setup/BUILD.md b/pydoc/setup/BUILD.md index 219d015f3a1..c18641133a8 100644 --- a/pydoc/setup/BUILD.md +++ b/pydoc/setup/BUILD.md @@ -110,6 +110,32 @@ conda install boost-cpp snappy thrift-cpp re2 utf8proc pip install boost snappy thrift re2 utf8proc ``` +#### Chapel Parquet Package + +The server build downloads and compiles the Chapel `Parquet` package into +`dep/build/Parquet`. Arrow and Parquet C++ libraries must still be available +through `pkg-config` as described above. + +You can optionally build against an existing package checkout or select a different +repository reference. You can also download the `Parquet` source ahead of time +using `wget https://github.com/chapel-lang/Parquet/archive/refs/heads/main.tar.gz`, +untar it, and then point `ARKOUDA_PARQUET_SRC_DIR` to the untarred location before running `make`. + +```bash +# Use an existing checkout without cloning it into dep/build. +ARKOUDA_PARQUET_SRC_DIR=/path/to/Parquet make + +# Select a branch or tag when creating a new checkout. +ARKOUDA_PARQUET_REF=my-branch make + +# Relocate Arkouda's downloaded checkout. +make PARQUET_INSTALL_DIR=/path/to/build/Parquet +``` + +Delete the downloaded checkout before changing `ARKOUDA_PARQUET_REF`; an +existing checkout is reused as-is. `ARKOUDA_PARQUET_REPO` can override the +default package repository. + #### Distributable Package Alternatively you can build a distributable package: diff --git a/scripts/get_parquet_package.sh b/scripts/get_parquet_package.sh new file mode 100755 index 00000000000..4dfc262e46e --- /dev/null +++ b/scripts/get_parquet_package.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# +# Fetch and build the external Chapel `Parquet` Mason package, then print the +# `chpl` arguments (C++ prerequisite headers/objects, include paths, Arrow link +# flags) and the Parquet module source path needed to compile it into the +# Arkouda server. +# +# This is used by the Makefile in place of `mason modules`, which does not work +# reliably with Chapel 2.4. +# +# Contract: ALL human-readable status is written to stderr. The ONLY thing +# written to stdout is a single line of `chpl` arguments, so a Make recipe can +# capture it with shell command substitution. +# +# Usage: +# get_parquet_package.sh +# +# Environment overrides: +# ARKOUDA_PARQUET_REPO git URL to clone (default: chapel-lang/Parquet) +# ARKOUDA_PARQUET_REF branch or tag to check out (default: repo HEAD) +# ARKOUDA_PARQUET_SRC_DIR use an existing checkout instead of cloning +# ARKOUDA_CHPL_HOME used to select Chapel's C++ compiler for prereqs + +set -euo pipefail + +log() { echo "get_parquet_package: $*" >&2; } + +PARQUET_REPO="${ARKOUDA_PARQUET_REPO:-https://github.com/chapel-lang/Parquet}" +PARQUET_REF="${ARKOUDA_PARQUET_REF:-}" + +INSTALL_DIR="${1:-${ARKOUDA_PARQUET_INSTALL_DIR:-}}" +if [[ -z "${INSTALL_DIR}" ]]; then + log "ERROR: no install directory provided (pass it as the first argument)" + exit 1 +fi + +# Allow pointing at an existing checkout (e.g. a Mason clone) to skip cloning. +PARQUET_SRC="${ARKOUDA_PARQUET_SRC_DIR:-${INSTALL_DIR}}" + +if [[ -f "${PARQUET_SRC}/Mason.toml" ]]; then + log "Using existing Parquet checkout at ${PARQUET_SRC}" +else + log "Cloning ${PARQUET_REPO}${PARQUET_REF:+@${PARQUET_REF}} into ${PARQUET_SRC}" + mkdir -p "$(dirname "${PARQUET_SRC}")" + git clone --depth 1 ${PARQUET_REF:+--branch "${PARQUET_REF}"} \ + "${PARQUET_REPO}" "${PARQUET_SRC}" >&2 +fi + +# Resolve to an absolute path so the emitted flags work from any CWD. +PARQUET_ROOT="$(cd "${PARQUET_SRC}" && pwd -P)" +PREREQ_DIR="${PARQUET_ROOT}/prereqs/cpp" +PARQUET_MODULE="${PARQUET_ROOT}/src/Parquet.chpl" + +if [[ ! -d "${PREREQ_DIR}" ]]; then + log "ERROR: expected C++ prerequisites at ${PREREQ_DIR}, but they are missing" + exit 1 +fi + +if [[ ! -f "${PARQUET_MODULE}" ]]; then + log "ERROR: expected Chapel module at ${PARQUET_MODULE}, but it is missing" + exit 1 +fi + +log "Building C++ prerequisites in ${PREREQ_DIR}" +make -s -C "${PREREQ_DIR}" CHPL_HOME="${ARKOUDA_CHPL_HOME:-}" >&2 + +# Gather the chpl flags the package needs. +FLAGS="$(make -s -C "${PREREQ_DIR}" CHPL_HOME="${ARKOUDA_CHPL_HOME}" printchplflags)" +if [[ -z "${FLAGS}" ]]; then + log "ERROR: the Parquet prerequisite build returned no Chapel flags" + exit 1 +fi + +printf '%s %s\n' "${FLAGS}" "${PARQUET_MODULE}" diff --git a/src/CheckpointMsg.chpl b/src/CheckpointMsg.chpl index 5eb0ff05d51..ceca286b0f4 100644 --- a/src/CheckpointMsg.chpl +++ b/src/CheckpointMsg.chpl @@ -386,7 +386,7 @@ module CheckpointMsg { return; } - if hasPrimitiveElements(entry) { + if hasPrimitiveElements(entry) && entry.etype != uint(8) { saveSymEntryPrimitive(entry, name, path, mdName, mdWriter); } else if entry.etype == bigint { diff --git a/src/ParquetMsg.chpl b/src/ParquetMsg.chpl index 795a4221cd6..11362b0159b 100644 --- a/src/ParquetMsg.chpl +++ b/src/ParquetMsg.chpl @@ -11,7 +11,6 @@ module ParquetMsg { use MultiTypeSymEntry; use NumPyDType; use Sort; - use CommAggregation; use AryUtil; use CTypes; use Map; @@ -19,141 +18,36 @@ module ParquetMsg { use IOUtils; use ParquetSharedEnums; - enum CompressionType { - NONE=0, - SNAPPY=1, - GZIP=2, - BROTLI=3, - ZSTD=4, - LZ4=5 - }; - + public use Parquet only CompressionType, ArrowTypes, ParquetError, + TRUNCATE, APPEND, ROWGROUPS, + ARROWINT64, ARROWINT32, ARROWUINT64, ARROWUINT32, + ARROWBOOLEAN, ARROWSTRING, ARROWFLOAT, ARROWDOUBLE, + ARROWLIST, ARROWDECIMAL, ARROWERROR, + readFilesByName, readStrFilesByName, + readListFilesByName, readColumn, readAllCols, + calcListSizesandOffset, calcStrSizesAndOffset, + calcStrListSizesAndOffset, getSubdomains, + getNullIndices, getStrColSize, getStrListColSize, + getArrSize, typeFromCType, typeToCType, getArrType, + getListData, getNumCols, getAllTypes, + populateTagData, getDatasets, getByteLength, + getVersionInfo, + write1DDistArrayParquet as masonWrite1DDistArrayParquet, + writeStringsColumn, writeListColumn, + writeStrListColumn, filesExistForWrite, pqWriteOp; // Use reflection for error information import Reflection.{getModuleName as getM, getRoutineName as getR, getLineNumber as getL}; - require "ReadParquet.h"; - require "ReadParquet.o"; - require "WriteParquet.h"; - require "WriteParquet.o"; - require "UtilParquet.h"; - require "UtilParquet.o"; private config const logLevel = ServerConfig.logLevel; private config const logChannel = ServerConfig.logChannel; const pqLogger = new Logger(logLevel, logChannel); - config const TRUNCATE: int = 0; - config const APPEND: int = 1; - - private config const ROWGROUPS = 512*1024*1024 / numBytes(int); // 512 mb of int64 + // Undocumented for now, just for internal experiments private config const batchSize = getEnvInt("ARKOUDA_SERVER_PARQUET_BATCH_SIZE", 8192); - extern var ARROWINT64: c_int; - extern var ARROWINT32: c_int; - extern var ARROWUINT64: c_int; - extern var ARROWUINT32: c_int; - extern var ARROWBOOLEAN: c_int; - extern var ARROWSTRING: c_int; - extern var ARROWFLOAT: c_int; - extern var ARROWLIST: c_int; - extern var ARROWDOUBLE: c_int; - extern var ARROWERROR: c_int; - extern var ARROWDECIMAL: c_int; - - extern record MyByteArray { - var len: uint(32); - var ptr: c_ptr(uint(8)); - }; - - enum ArrowTypes { int64, int32, uint64, uint32, - stringArr, timestamp, boolean, - double, float, list, decimal, - notimplemented }; - - record parquetErrorMsg { - var errMsg: c_ptr(uint(8)); - proc init() { - errMsg = nil; - } - - proc deinit() { - extern proc c_free_string(ptr); - c_free_string(errMsg); - } - - proc parquetError(lineNumber, routineName, moduleName) throws { - extern proc strlen(a): int; - var err: string; - try { - err = string.createCopyingBuffer(errMsg, strlen(errMsg)); - } catch e { - err = "Error converting Parquet error message to Chapel string"; - } - if err == "Unexpected end of stream" { - err += ". This may be due to null values in a non-float column. Try again with the flag has_non_float_nulls=True"; - } - throw getErrorWithContext( - msg=err, - lineNumber, - routineName, - moduleName, - errorClass="ParquetError"); - } - } - - proc getVersionInfo() { - extern proc c_getVersionInfo(): c_ptrConst(c_char); - extern proc strlen(str): c_int; - extern proc c_free_string(ptr); - var cVersionString = c_getVersionInfo(); - defer { - c_free_string(cVersionString: c_ptr(void)); - } - var ret: string; - try { - ret = string.createCopyingBuffer(cVersionString, - strlen(cVersionString)); - } catch e { - ret = "Error converting Arrow version message to Chapel string"; - } - return ret; - } - - proc getSubdomains(lengths: [?FD] int) { - var subdoms: [FD] domain(1); - var offset = 0; - for i in FD { - subdoms[i] = {offset..#lengths[i]}; - offset += lengths[i]; - } - return subdoms; - } - - proc getRGSubdomains(bytesPerRG: [?D] ?t, maxRowGroups: int) { - var rgSubdomains: [D] [0..#maxRowGroups] domain(1); - - var offset = 0; - for i in D { - for rg in 0..#maxRowGroups { - if bytesPerRG[i][rg] != 0 { - rgSubdomains[i][rg] = {offset..#bytesPerRG[i][rg]}; - offset += bytesPerRG[i][rg]; - } - } - } - return (rgSubdomains, offset); - } - - inline proc readFilesByName(ref A: [] ?t, filenames: [] string, sizes: [] int, - dsetname: string, ty, byteLength=-1, - hasNonFloatNulls=false) throws { - var dummy = [false]; - readFilesByName(A, dummy, filenames, sizes, dsetname, ty, byteLength, - hasNonFloatNulls, hasWhereNull=false); - } - class ParquetReadError: ErrorWithContext { proc init(msg: string, moduleName: string, routineName: string, lineNumber: int(64)) { @@ -187,569 +81,20 @@ module ParquetMsg { } } - proc getCPtrToAkArrayEntry(e: borrowed GenSymEntry, - ctype: c_int, off: int): c_ptr(void) throws { - var ret: c_ptr(void); - select typeFromCType(ctype) { - when ArrowTypes.int64 do - ret = c_ptrTo((e:(borrowed SymEntry(int(64),1))).a[off]); - when ArrowTypes.uint64 do - ret = c_ptrTo((e:(borrowed SymEntry(uint(64),1))).a[off]); - when ArrowTypes.double do - ret = c_ptrTo((e:(borrowed SymEntry(real(64),1))).a[off]); - otherwise do - throw new NotImplementedError("Unexpected column type while " + - "reading parquet file", getL(), getR(), getM()); - } - return ret; - } - - - /* - whereNull will be populated by the CPP interface, where `true` would mean a - 0 (null) having been read. - */ - proc readFilesByName(ref A: [] ?t, ref whereNull: [] bool, - filenames: [] string, sizes: [] int, dsetname: string, - ty, byteLength=-1, hasNonFloatNulls=false, - param hasWhereNull=true) throws { - extern proc c_readColumnByName(filename, arr_chpl, where_null_chpl, colNum, numElems, startIdx, batchSize, byteLength, hasNonFloatNulls, errMsg): int; - - var subdoms = getSubdomains(sizes); - var fileOffsets = (+ scan sizes) - sizes; - - coforall loc in A.targetLocales() with (ref A) do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - var locOffsets = fileOffsets; - - forall (off, filedom, filename) in zip(locOffsets, locFiledoms, locFiles) { - for locdom in A.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - if intersection.size > 0 { - var pqErr = new parquetErrorMsg(); - var whereNullPtr = if hasWhereNull - then c_ptrTo(whereNull[intersection.low]) - else nil; - if c_readColumnByName(filename.localize().c_str(), - c_ptrTo(A[intersection.low]), - whereNullPtr, - dsetname.localize().c_str(), - intersection.size, intersection.low - off, - batchSize, byteLength, hasNonFloatNulls, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - } - } - } - } - - proc readStrFilesByName(ref A: [] ?t, filenames: [] string, sizes: [] int, dsetname: string) throws { - extern proc c_readStrColumnByName(filename, arr_chpl, colname, numElems, batchSize, errMsg): int; - var subdoms = getSubdomains(sizes); - - coforall loc in A.targetLocales() do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - - forall (filedom, filename) in zip(locFiledoms, locFiles) { - for locdom in A.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - - if intersection.size > 0 { - var pqErr = new parquetErrorMsg(); - var col: [filedom] t; - - if c_readStrColumnByName(filename.localize().c_str(), c_ptrTo(col), - dsetname.localize().c_str(), filedom.size, - batchSize, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - A[filedom] = col; - } - } - } - } - } - - proc readListFilesByName(A: [] ?t, rows_per_file: [] int, seg_sizes: [] int, offsets: [] int, filenames: [] string, sizes: [] int, dsetname: string, ty) throws { - extern proc c_readListColumnByName(filename, arr_chpl, colNum, numElems, startIdx, batchSize, errMsg): int; - var subdoms = getSubdomains(sizes); - var fileOffsets = (+ scan sizes) - sizes; - var segmentOffsets = (+ scan rows_per_file) - rows_per_file; - - coforall loc in A.targetLocales() do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - var locOffsets = fileOffsets; // value count offset - var locSegOffsets = segmentOffsets; // indicates which segment index is first for the file - - forall (s, off, filedom, filename) in zip(locSegOffsets, locOffsets, locFiledoms, locFiles) { - for locdom in A.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - - if intersection.size > 0 { - var pqErr = new parquetErrorMsg(); - var col: [filedom] t; - if c_readListColumnByName(filename.localize().c_str(), c_ptrTo(col), - dsetname.localize().c_str(), filedom.size, 0, - batchSize, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - A[filedom] = col; - } - } - } - } - } - - proc calcListSizesandOffset(seg_sizes: [] ?t, filenames: [] string, sizes: [] int, dsetname: string) throws { - var subdoms = getSubdomains(sizes); - - var listSizes: [filenames.domain] int; - var file_offset: int = 0; - coforall loc in seg_sizes.targetLocales() with (ref listSizes) do on loc{ - var locFiles = filenames; - var locFiledoms = subdoms; - - forall (i, filedom, filename) in zip(sizes.domain, locFiledoms, locFiles) { - for locdom in seg_sizes.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - if intersection.size > 0 { - var col: [filedom] t; - listSizes[i] = getListColSize(filename, dsetname, col); - seg_sizes[filedom] = col; // this is actually segment sizes here - } - } - } - } - return listSizes; - } - - proc calcStrSizesAndOffset(offsets: [] ?t, filenames: [] string, sizes: [] int, dsetname: string) throws { - var subdoms = getSubdomains(sizes); - - var byteSizes: [filenames.domain] int; - - coforall loc in offsets.targetLocales() with (ref byteSizes) do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - - forall (i, filedom, filename) in zip(sizes.domain, locFiledoms, locFiles) { - for locdom in offsets.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - if intersection.size > 0 { - var col: [filedom] t; - byteSizes[i] = getStrColSize(filename, dsetname, col); - offsets[filedom] = col; - } - } - } - } - return byteSizes; - } - - proc calcStrListSizesAndOffset(offsets: [] ?t, filenames: [] string, sizes: [] int, dsetname: string) throws { - var subdoms = getSubdomains(sizes); - - var byteSizes: [filenames.domain] int; - - coforall loc in offsets.targetLocales() with (ref byteSizes) do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - - forall (i, filedom, filename) in zip(sizes.domain, locFiledoms, locFiles) { - for locdom in offsets.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - if intersection.size > 0 { - var col: [filedom] t; - byteSizes[i] = getStrListColSize(filename, dsetname, col); - offsets[filedom] = col; - } - } - } - } - return byteSizes; - } - - proc getNullIndices(A: [] ?t, filenames: [] string, sizes: [] int, dsetname: string, ty) throws { - extern proc c_getStringColumnNullIndices(filename, colname, nulls_chpl, errMsg): int; - var subdoms = getSubdomains(sizes); - - coforall loc in A.targetLocales() do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - - forall (filedom, filename) in zip(locFiledoms, locFiles) { - for locdom in A.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - - if intersection.size > 0 { - var pqErr = new parquetErrorMsg(); - var col: [filedom] t; - if c_getStringColumnNullIndices(filename.localize().c_str(), dsetname.localize().c_str(), - c_ptrTo(col), pqErr.errMsg) { - pqErr.parquetError(getL(), getR(), getM()); - } - A[filedom] = col; - } - } - } - } - } - - proc getStrColSize(filename: string, dsetname: string, ref offsets: [] int) throws { - 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, 256, - c_ptrTo(pqErr.errMsg)); - - if byteSize == ARROWERROR then - pqErr.parquetError(getL(), getR(), getM()); - return byteSize; - } - - proc getStrListColSize(filename: string, dsetname: string, ref offsets: [] int) throws { - extern proc c_getStringListColumnNumBytes(filename, colname, offsets, numElems, startIdx, batchSize, errMsg): int; - var pqErr = new parquetErrorMsg(); - - var byteSize = c_getStringListColumnNumBytes(filename.localize().c_str(), - dsetname.localize().c_str(), - c_ptrTo(offsets), - offsets.size, 0, 256, - c_ptrTo(pqErr.errMsg)); - - if byteSize == ARROWERROR then - pqErr.parquetError(getL(), getR(), getM()); - return byteSize; - } - - proc getListColSize(filename: string, dsetname: string, ref seg_sizes: [] int) throws { - extern proc c_getListColumnSize(filename, colname, seg_sizes, numElems, startIdx, errMsg): int; - var pqErr = new parquetErrorMsg(); - - var listSize = c_getListColumnSize(filename.localize().c_str(), - dsetname.localize().c_str(), - c_ptrTo(seg_sizes), - seg_sizes.size, 0, - c_ptrTo(pqErr.errMsg)); - - if listSize == ARROWERROR then - pqErr.parquetError(getL(), getR(), getM()); - return listSize; - } - - proc getArrSize(filename: string) throws { - extern proc c_getNumRows(str_chpl, errMsg): int; - var pqErr = new parquetErrorMsg(); - - var size = c_getNumRows(filename.localize().c_str(), - c_ptrTo(pqErr.errMsg)); - if size == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - return size; - } - - proc typeFromCType(ctype: c_int) throws { - if ctype == ARROWINT64 then return ArrowTypes.int64; - else if ctype == ARROWINT32 then return ArrowTypes.int32; - else if ctype == ARROWUINT32 then return ArrowTypes.uint32; - else if ctype == ARROWUINT64 then return ArrowTypes.uint64; - else if ctype == ARROWBOOLEAN then return ArrowTypes.boolean; - else if ctype == ARROWSTRING then return ArrowTypes.stringArr; - else if ctype == ARROWDOUBLE then return ArrowTypes.double; - else if ctype == ARROWFLOAT then return ArrowTypes.float; - else if ctype == ARROWLIST then return ArrowTypes.list; - else if ctype == ARROWDECIMAL then return ArrowTypes.decimal; - throw getErrorWithContext(getL(), getM(), getR(), - msg="Unrecognized Parquet data type", - errorClass="ParquetError"); - return ArrowTypes.notimplemented; - } - - proc ctypeFromType(t: ArrowTypes) throws { - if t == ArrowTypes.int64 then return ARROWINT64; - else if t == ArrowTypes.int32 then return ARROWINT32; - else if t == ArrowTypes.uint32 then return ARROWUINT32; - else if t == ArrowTypes.uint64 then return ARROWUINT64; - else if t == ArrowTypes.boolean then return ARROWBOOLEAN; - else if t == ArrowTypes.stringArr then return ARROWSTRING; - else if t == ArrowTypes.double then return ARROWDOUBLE; - else if t == ArrowTypes.float then return ARROWFLOAT; - else if t == ArrowTypes.list then return ARROWLIST; - else if t == ArrowTypes.decimal then return ARROWDECIMAL; - throw getErrorWithContext(getL(), getM(), getR(), - msg="Unrecognized Parquet data type", - errorClass="ParquetError"); - return ARROWERROR; - } - - proc getArrType(filename: string, colname: string) throws { - extern proc c_getType(filename, colname, errMsg): c_int; - var pqErr = new parquetErrorMsg(); - var arrType = c_getType(filename.localize().c_str(), - colname.localize().c_str(), - c_ptrTo(pqErr.errMsg)); - - if arrType == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - - return typeFromCType(arrType); - } - - proc getListData(filename: string, dsetname: string) throws { - extern proc c_getListType(filename, dsetname, errMsg): c_int; - var pqErr = new parquetErrorMsg(); - - var t = c_getListType(filename.localize().c_str(), dsetname.localize().c_str(), c_ptrTo(pqErr.errMsg)); - if t == ARROWINT64 then return ArrowTypes.int64; - else if t == ARROWINT32 then return ArrowTypes.int32; - else if t == ARROWUINT32 then return ArrowTypes.uint32; - else if t == ARROWUINT64 then return ArrowTypes.uint64; - else if t == ARROWBOOLEAN then return ArrowTypes.boolean; - else if t == ARROWSTRING then return ArrowTypes.stringArr; - else if t == ARROWDOUBLE then return ArrowTypes.double; - else if t == ARROWFLOAT then return ArrowTypes.float; - return ArrowTypes.notimplemented; - } - - proc toCDtype(dtype: string) throws { - select dtype { - when 'int64' { - return ARROWINT64; - } when 'uint32' { - return ARROWUINT32; - } when 'uint64' { - return ARROWUINT64; - } when 'bool' { - return ARROWBOOLEAN; - } when 'float64' { - return ARROWDOUBLE; - } when 'str' { - return ARROWSTRING; - } otherwise { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Trying to convert unrecognized dtype to Parquet type", - errorClass="ParquetError"); - return ARROWERROR; - } - } - } - - proc writeDistArrayToParquet(A, filename, dsetname, dtype, rowGroupSize, compression, mode) throws { - extern proc c_writeColumnToParquet(filename, arr_chpl, colnum, - dsetname, numelems, rowGroupSize, - dtype, compression, errMsg): int; - extern proc c_appendColumnToParquet(filename, arr_chpl, - dsetname, numelems, - dtype, compression, - errMsg): int; - var dtypeRep = toCDtype(dtype); - var prefix: string; - var extension: string; - - (prefix, extension) = getFileMetadata(filename); - - // Generate the filenames based upon the number of targetLocales. - var filenames = generateFilenames(prefix, extension, A.targetLocales().size); - var numElemsPerFile: [filenames.domain] int; - - //Generate a list of matching filenames to test against. - var matchingFilenames = getMatchingFilenames(prefix, extension); - - var filesExist = processParquetFilenames(filenames, matchingFilenames, mode); - - if mode == APPEND { - if filesExist { - var datasets = getDatasets(filenames[0]); - if datasets.contains(dsetname) then - throw getErrorWithContext(getL(), getM(), getR(), - msg="A column with name " + dsetname + - " already exists in Parquet file", - errorClass='WriteModeError'); - } - } - - coforall (loc, idx) in zip(A.targetLocales(), filenames.domain) do on loc { - var pqErr = new parquetErrorMsg(); - const myFilename = filenames[idx]; - - var locDom = A.localSubdomain(); - var locArr = A[locDom]; - - numElemsPerFile[idx] = locDom.size; - - var valPtr: c_ptr(void) = nil; - if locArr.size != 0 { - 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(getL(), getR(), getM()); - } - } else { - if c_appendColumnToParquet(myFilename.localize().c_str(), valPtr, - dsetname.localize().c_str(), locDom.size, - dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - } - // Only warn when files are being overwritten in truncate mode - return (filesExist && mode == TRUNCATE, filenames, numElemsPerFile); - } - - proc createEmptyParquetFile(filename: string, dsetname: string, dtype: int, compression: int) throws { - extern proc c_createEmptyParquetFile(filename, dsetname, dtype, - compression, errMsg): int; - var pqErr = new parquetErrorMsg(); - if c_createEmptyParquetFile(filename.localize().c_str(), dsetname.localize().c_str(), - dtype, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - // TODO: do we want to add offset writing for Parquet string writes? // if we do, then we need to add the load offsets functionality // in the string reading function - proc write1DDistStringsAggregators(filename: string, mode: int, dsetName: string, entry: SegStringSymEntry, compression: int) throws { + proc write1DDistStrings(filename: string, mode: int, dsetName: string, + entry: SegStringSymEntry, compression: int) throws { var segString = new SegString("", entry); ref ss = segString; - var A = ss.offsets.a; - - var prefix: string; - var extension: string; - - (prefix, extension) = getFileMetadata(filename); - - // Generate the filenames based upon the number of targetLocales. - var filenames = generateFilenames(prefix, extension, A.targetLocales().size); - - //Generate a list of matching filenames to test against. - var matchingFilenames = getMatchingFilenames(prefix, extension); - - var filesExist = processParquetFilenames(filenames, matchingFilenames, mode); - - if mode == APPEND { - if filesExist { - var datasets = getDatasets(filenames[0]); - if datasets.contains(dsetName) then - throw getErrorWithContext(getL(), getM(), getR(), - msg="A column with name " + dsetName + - " already exists in Parquet file", - errorClass='WriteModeError'); - } - } - - const extraOffset = ss.values.size; - const lastOffset = if A.size == 0 then 0 else A[A.domain.high]; // prevent index error when empty - const lastValIdx = ss.values.a.domain.high; - // For each locale gather the string bytes corresponding to the offsets in its local domain - coforall (loc, idx) in zip(A.targetLocales(), filenames.domain) with (ref ss) do on loc { - const myFilename = filenames[idx]; - - const locDom = A.localSubdomain(); - var dims: [0..#1] int; - dims[0] = locDom.size: int; - - if (locDom.isEmpty() || locDom.size <= 0) { - if mode == APPEND && filesExist then - throw getErrorWithContext(getL(), getM(), getR(), - msg="Parquet columns must each have the same length: " + - myFilename, - errorClass='WriteModeError'); - createEmptyParquetFile(myFilename, dsetName, ARROWSTRING, compression); - } else { - var localOffsets = A[locDom]; - var startValIdx = localOffsets[locDom.low]; - - var endValIdx = if (lastOffset == localOffsets[locDom.high]) then lastValIdx else A[locDom.high + 1] - 1; - - var valIdxRange = startValIdx..endValIdx; - var localVals: [valIdxRange] uint(8); - ref olda = ss.values.a; - forall (localVal, valIdx) in zip(localVals, valIdxRange) with (var agg = newSrcAggregator(uint(8))) { - // Copy the remote value at index position valIdx to our local array - agg.copy(localVal, olda[valIdx]); // in SrcAgg, the Right Hand Side is REMOTE - } - var locOffsets: [0..#locDom.size+1] int; - locOffsets[0..#locDom.size] = A[locDom]; - if locDom.high == A.domain.high then - locOffsets[locOffsets.domain.high] = extraOffset; - else - locOffsets[locOffsets.domain.high] = A[locDom.high+1]; - - writeStringsComponentToParquet(myFilename, dsetName, localVals, locOffsets, ROWGROUPS, compression, mode, filesExist); - } - } - return filesExist && mode == TRUNCATE; - } - - private proc writeStringsComponentToParquet(filename, dsetname, ref values: [] uint(8), ref offsets: [] int, rowGroupSize, compression, mode, filesExist) throws { - extern proc c_writeStrColumnToParquet(filename, arr_chpl, offsets_chpl, - dsetname, numelems, rowGroupSize, - dtype, compression, errMsg): int; - extern proc c_appendColumnToParquet(filename, arr_chpl, - dsetname, numelems, - dtype, compression, - errMsg): int; - var pqErr = new parquetErrorMsg(); - var dtypeRep = ARROWSTRING; - if mode == TRUNCATE || !filesExist { - if c_writeStrColumnToParquet(filename.localize().c_str(), c_ptrTo(values), c_ptrTo(offsets), - dsetname.localize().c_str(), offsets.size-1, rowGroupSize, - dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } else if mode == APPEND { - if c_appendColumnToParquet(filename.localize().c_str(), c_ptrTo(values), - dsetname.localize().c_str(), offsets.size-1, - dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - } - - proc processParquetFilenames(filenames: [] string, matchingFilenames: [] string, mode: int) throws { - var filesExist: bool = true; - if mode == APPEND { - if matchingFilenames.size == 0 { - // Files do not exist, so we can just create the files - filesExist = false; - } - else if matchingFilenames.size != filenames.size { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Appending to existing files must be done with the same number " + - "of locales. Try saving with a different directory or filename prefix?", - errorClass='MismatchedAppendError' - ); - } - } else if mode == TRUNCATE { - if matchingFilenames.size > 0 { - filesExist = true; - } else { - filesExist = false; - } - } else { - throw getErrorWithContext(getL(), getM(), getR(), - msg="The mode %? is invalid".format(mode), - errorClass='IllegalArgumentError'); - } - return filesExist; + return writeStringsColumn(filename, dsetName, ss.offsets.a, ss.values.a, + compression: CompressionType, mode); } proc write1DDistArrayParquet(filename: string, dsetname, dtype, compression, mode, A) throws { - return writeDistArrayToParquet(A, filename, dsetname, dtype, ROWGROUPS, compression, mode); + return masonWrite1DDistArrayParquet(filename, dsetname, + compression: CompressionType, mode, A); } proc parseListDataset(filenames: [] string, dsetname: string, ty, len: int, sizes: [] int, st: borrowed SymTab) throws { @@ -805,126 +150,6 @@ module ParquetMsg { return formatJson(rtnmap); } - proc populateTagData(A, filenames: [?fD] string, sizes) throws { - var subdoms = getSubdomains(sizes); - var fileOffsets = (+ scan sizes) - sizes; - - coforall loc in A.targetLocales() do on loc { - var locFiles = filenames; - var locFiledoms = subdoms; - var locOffsets = fileOffsets; - - try { - forall (off, filedom, filename, tag) in zip(locOffsets, locFiledoms, locFiles, 0..) { - for locdom in A.localSubdomains() { - const intersection = domain_intersection(locdom, filedom); - - if intersection.size > 0 { - // write the tag into the entry - A[intersection] = tag; - } - } - } - } - } - } - - inline proc getReaderIdx(fileNum: int, rgNum: int) { - // we can assume there won't be more than 1000 RGs in a file - return (fileNum*1000) + rgNum; - } - - proc getRowGroupNums(ref distFiles, ref numRowGroups) { - coforall loc in distFiles.targetLocales() with (ref numRowGroups) do on loc { - var locFiles: [distFiles.localSubdomain()] string = distFiles[distFiles.localSubdomain()]; - for i in locFiles.domain { - c_openFile(locFiles[i].localize().c_str(), getReaderIdx(i,0)); - numRowGroups[i] = c_getNumRowGroups(getReaderIdx(i,0)); - for j in 2..numRowGroups[i] { - c_openFile(locFiles[i].localize().c_str(), getReaderIdx(i,j-1)); - } - } - } - var maxRowGroups = 0; - for val in numRowGroups do if maxRowGroups < val then maxRowGroups = val; - return maxRowGroups; - } - - proc fillSegmentsAndPersistData(ref distFiles, ref entrySeg, ref externalData, ref containsNulls, ref valsRead, dsetname, sizes, len, numRowGroups, ref bytesPerRG, ref startIdxs) throws { - var subdoms = getSubdomains(sizes); - coforall loc in distFiles.targetLocales() with (ref externalData, ref valsRead, ref bytesPerRG) do on loc { - var locFiles: [distFiles.localSubdomain()] string = distFiles[distFiles.localSubdomain()]; - var locSubdoms = subdoms; - - for i in locFiles.domain { - var fname = locFiles[i]; - var locDsetname = dsetname; - for rg in 0..#numRowGroups[i] { - c_createRowGroupReader(rg, getReaderIdx(i,rg)); - c_createColumnReader(locDsetname.localize().c_str(), getReaderIdx(i,rg)); - } - } - - var errs: [locFiles.domain] (bool, parquetErrorMsg) = [0..#locFiles.domain] (false, new parquetErrorMsg()); - forall i in locFiles.domain { - var fname = locFiles[i]; - var locDsetname = dsetname; - var startIdx = locSubdoms[i].low; - for rg in 0..#numRowGroups[i] { - var totalBytes = 0; - startIdxs[i][rg] = startIdx; - - var numRead = 0; - - if c_readParquetColumnChunks(fname.localize().c_str(), batchSize, len, getReaderIdx(i,rg), c_ptrTo(numRead), c_ptrTo(externalData[i][rg]), c_ptrTo(containsNulls[i][rg]), c_ptrTo(errs[i][1].errMsg)) == ARROWERROR { - errs[i] = (true, errs[i][1]); - } - var tmp: [startIdx..#numRead] int; - forall (id, j) in zip(0..#numRead, startIdx..#numRead) with (+ reduce totalBytes) { - ref curr = (externalData[i][rg]: c_ptr(MyByteArray))[id]; - tmp[j] = curr.len + 1; // this was only change - totalBytes += curr.len+1; - } - entrySeg.a[startIdx..#numRead] = tmp; - valsRead[i][rg] = numRead; - startIdx += numRead; - bytesPerRG[i][rg] = totalBytes; - totalBytes = 0; - } - } - for (hadErr, err) in errs do - if hadErr then - err.parquetError(getL(), getR(), getM()); - } - } - - proc copyValuesFromC(ref entryVal, ref distFiles, ref externalData, ref valsRead, ref numRowGroups, ref rgSubdomains, maxRowGroups, sizes, ref segArr, ref startIdxs) { - var subdoms = getSubdomains(sizes); - coforall loc in distFiles.targetLocales() with (ref externalData) do on loc { - var locValsRead: [valsRead.localSubdomain()] [0..#maxRowGroups] int = valsRead[valsRead.localSubdomain()]; - var locNumRowGroups: [numRowGroups.localSubdomain()] int = numRowGroups[numRowGroups.localSubdomain()]; - var locStartIdxs: [startIdxs.localSubdomain()] [0..#maxRowGroups] int = startIdxs[startIdxs.localSubdomain()]; - var locSubdoms = subdoms; - - forall i in locNumRowGroups.domain { - var numRgs = locNumRowGroups[i]; - for rg in 0..#numRgs { - var entryIdx = rgSubdomains[i][rg].low; - var numRead = locValsRead[i][rg]; - var offsetIdx = locStartIdxs[i][rg]; - var tmp: [rgSubdomains[i][rg]] uint(8); - forall (idx, oIdx) in zip(0..#numRead, offsetIdx..#numRead) { - ref curr = (externalData[i][rg]: c_ptr(MyByteArray))[idx]; - for j in 0..#curr.len { - tmp[segArr[oIdx]+j] = curr.ptr[j]; - } - } - entryVal.a[rgSubdomains[i][rg]] = tmp; - } - } - } - } - proc readAllColsParquetMsg(cmd: string, msgArgs: borrowed MessageArgs, st: borrowed SymTab): MsgTuple throws { var repMsg: string; @@ -1014,21 +239,10 @@ module ParquetMsg { } } - var pqErr = new parquetErrorMsg(); - - extern proc c_getNumCols(filename, errMsg): int(64); - const numCols = c_getNumCols(Filenames[0].c_str(), c_ptrTo(pqErr.errMsg)); - if numCols == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } + const numCols = getNumCols(Filenames[0]); const ColDomain = {0..#numCols}; - var CTypes: [ColDomain] c_int; - - extern proc c_getAllTypes(filename, types_out, errMsg): c_int; - if c_getAllTypes(Filenames[0].c_str(), c_ptrTo(CTypes), c_ptrTo(pqErr.errMsg)) { - pqErr.parquetError(getL(), getR(), getM()); - } + const CTypes: [ColDomain] c_int = getAllTypes(Filenames[0]); var op = new pqReadColOp(Filenames, len, CTypes, Sizes, Offsets, hasNonFloatNulls, nullMode); @@ -1220,10 +434,6 @@ module ParquetMsg { proc ref readInto(ref e: shared GenSymEntry?) throws where this.kind == pqReadReqKind.oneCol { - extern proc c_readColumnByName(filename, arr_chpl, where_null_chpl, - colName, numElems, startIdx, batchSize, - byteLength, hasNonFloatNulls, errMsg): int; - // TODO in the all-col implementation, this is handled at CPP const byteLength = if ty == ARROWDECIMAL then getByteLength(filenames[0], dsetname) @@ -1253,20 +463,18 @@ module ParquetMsg { for locdom in Dom.localSubdomains() { const intersection = domain_intersection(locdom, filedom); if intersection.size > 0 { - var pqErr = new parquetErrorMsg(); var whereNullPtr = if hasNonFloatNulls then c_ptrTo(whereNull[intersection.low]) else nil; - - if c_readColumnByName(filename.localize().c_str(), - getPtr(e, intersection.low), - whereNullPtr, - dsetname.localize().c_str(), - intersection.size, intersection.low - off, - batchSize, byteLength, hasNonFloatNulls, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } + readColumn(filename=filename, + colName=dsetname, + ptr=getPtr(e, intersection.low), + whereNullPtr=whereNullPtr, + numElems=intersection.size, + startIdx=intersection.low - off, + batchSize=batchSize, + byteLength=byteLength, + hasNonFloatNulls=hasNonFloatNulls); } } } @@ -1277,10 +485,6 @@ module ParquetMsg { proc ref readInto(ref e: [] shared GenSymEntry?) throws where this.kind == pqReadReqKind.allCols { - extern proc c_readAllCols(filename, chpl_arrs, types, where_null_chpl, - numElems, startIdx, batchSize, - nullMode, errMsg): c_int; - var subdoms = getSubdomains(this.sizes); pqLogger.debug(getM(), getR(), getL(), @@ -1313,19 +517,15 @@ module ParquetMsg { } - var pqErr = new parquetErrorMsg(); const startIdx = intersection.low - off; pqLogger.debug(getM(), getR(), getL(), "about to call c_readAllCols for locale %? with filename=%?, startIdx=%?, ptrs=%?" .format(here.id, filename, startIdx, CPtrsToData)); - if c_readAllCols(filename.localize().c_str(), - c_ptrTo(CPtrsToData), c_ptrToConst(LocTypes), - c_ptrTo(CPtrsToWhereNulls), - numElems=intersection.size, - startIdx=intersection.low-off, - batchSize, nullMode, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } + readAllCols(filename, CPtrsToData, LocTypes, + CPtrsToWhereNulls, + numElems=intersection.size, + startIdx=intersection.low-off, + batchSize=batchSize, + nullMode=nullMode:int); } } } @@ -1415,7 +615,10 @@ module ParquetMsg { validUIntType(t) || t==bool { const ref myWhereNull = _whereNull[colIdx]; - if hasNonFloatNulls && (|| reduce myWhereNull) { + const handleNonFloatNulls = if isOneCol() + then hasNonFloatNulls + else nullMode == NullMode.all; + if handleNonFloatNulls && (|| reduce myWhereNull) { // if we have non-float nulls and there's at least one null var floatEntry = createSymEntry(se.size, real); floatEntry.a = (se.a):real; @@ -1554,7 +757,7 @@ module ParquetMsg { tagData = false; // turn off so we only run once } - const cty = ctypeFromType(ty); + const cty = typeToCType(ty); if pqReadColOp.canHandleType(cty) { var dummyNullMode: NullMode; // ignored for single col reads for now @@ -1589,17 +792,20 @@ module ParquetMsg { "%s+%?".format(stringsEntry.name, stringsEntry.nBytes))); } else if ty == ArrowTypes.list { - var list_ty = getListData(filenames[0], dsetname); + var list_ty: ArrowTypes; + try { + list_ty = getListData(filenames[0], dsetname); // check for and skip further nested datasets - if list_ty == ArrowTypes.notimplemented { + } catch e: ParquetError { pqLogger.info(getM(),getR(),getL(), - "Invalid list datatype found in %s. Skipping.".format(dsetname)); - } - else { - var create_str: string = parseListDataset(filenames, dsetname, - list_ty, len, sizes, st); - rnames.pushBack((dsetname, ObjType.SEGARRAY, create_str)); + "Invalid list datatype found in %s (%s). Skipping." + .format(dsetname, e.message())); + continue; } + + var create_str: string = parseListDataset(filenames, dsetname, + list_ty, len, sizes, st); + rnames.pushBack((dsetname, ObjType.SEGARRAY, create_str)); } else { var errorMsg = "DType %s not supported for Parquet reading".format(ty); pqLogger.error(getM(),getR(),getL(),errorMsg); @@ -1612,76 +818,6 @@ module ParquetMsg { return new MsgTuple(repMsg,MsgType.NORMAL); } - iter datasets(filename) { - extern proc c_getDatasetNames(filename, dsetResult, readNested, errMsg): int(32); - extern proc strlen(a): int; - var pqErr = new parquetErrorMsg(); - var res: c_ptr(uint(8)); - defer { - extern proc c_free_string(ptr); - c_free_string(res); - } - if c_getDatasetNames(filename.c_str(), c_ptrTo(res), false, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - var datasets: string; - try! datasets = string.createCopyingBuffer(res, strlen(res)); - for s in datasets.split(",") do yield s; - } - - // TODO remove this and use the iterator everywhere, or turn this into a - // list-returning version - proc getDatasets(filename) throws { - extern proc c_getDatasetNames(filename, dsetResult, readNested, errMsg): int(32); - extern proc strlen(a): int; - var pqErr = new parquetErrorMsg(); - var res: c_ptr(uint(8)); - defer { - extern proc c_free_string(ptr); - c_free_string(res); - } - if c_getDatasetNames(filename.c_str(), c_ptrTo(res), false, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - var datasets: string; - datasets = string.createCopyingBuffer(res, strlen(res)); - return new list(datasets.split(",")); - } - - // Decimal columns in Parquet have a fixed number of bytes based on the precision, - // but there isn't a way in Parquet to get the precision. Since the byte length - // will always remain the same for each precision value, here we just created a - // lookup table that maps from the precision to the byte value. - proc getByteLength(filename, colname) throws { - extern proc c_getPrecision(filename, colname, errMsg): int(32); - var pqErr = new parquetErrorMsg(); - var res: c_ptr(uint(8)); - defer { - extern proc c_free_string(ptr); - c_free_string(res); - } - - var precision = c_getPrecision(filename.c_str(), colname.c_str(), c_ptrTo(pqErr.errMsg)); - if precision < 3 then return 1; - else if precision < 5 then return 2; - else if precision < 7 then return 3; - else if precision < 10 then return 4; - else if precision < 12 then return 5; - else if precision < 15 then return 6; - else if precision < 17 then return 7; - else if precision < 19 then return 8; - else if precision < 22 then return 9; - else if precision < 24 then return 10; - else if precision < 27 then return 11; - else if precision < 29 then return 12; - else if precision < 32 then return 13; - else if precision < 34 then return 14; - else if precision < 36 then return 15; - return 16; - } - proc pdarray_toParquetMsg(msgArgs: MessageArgs, st: borrowed SymTab): bool throws { var mode = msgArgs.get("mode").getIntValue(); var filename: string = msgArgs.getValueOf("prefix"); @@ -1742,192 +878,20 @@ module ParquetMsg { } var segString:SegStringSymEntry = toSegStringSymEntry(entry); - var warnFlag: bool = write1DDistStringsAggregators(filename, mode, dsetname, segString, compression:int); + var warnFlag: bool = write1DDistStrings(filename, mode, dsetname, segString, compression:int); return warnFlag; } - proc createEmptyListParquetFile(filename: string, dsetname: string, dtype: int, compression: int) throws { - extern proc c_createEmptyListParquetFile(filename, dsetname, dtype, - compression, errMsg): int; - var pqErr = new parquetErrorMsg(); - if c_createEmptyListParquetFile(filename.localize().c_str(), dsetname.localize().c_str(), - dtype, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - - proc writeSegArrayComponent(filename: string, dsetname: string, const ref distVals: [] ?t, valIdxRange, segments, locDom, - extraOffset, lastOffset, lastValId, c_dtype, compression) throws { - extern proc c_writeListColumnToParquet(filename, arr_chpl, offsets_chpl, - dsetname, numelems, rowGroupSize, - dtype, compression, errMsg): int; - var localVals: [valIdxRange] t = distVals[valIdxRange]; - var locOffsets: [0..#locDom.size+1] int; - locOffsets[0..#locDom.size] = segments[locDom]; - if locDom.high == segments.domain.high then - locOffsets[locOffsets.domain.high] = extraOffset; - else - locOffsets[locOffsets.domain.high] = segments[locDom.high+1]; - - var pqErr = new parquetErrorMsg(); - - var valPtr: c_ptr(void) = nil; - if localVals.size != 0 { - valPtr = c_ptrTo(localVals); - } - - if c_writeListColumnToParquet(filename.localize().c_str(), c_ptrTo(locOffsets), valPtr, - dsetname.localize().c_str(), locOffsets.size-1, ROWGROUPS, - c_dtype, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - proc writeSegArrayParquet(filename: string, dsetName: string, c_dtype, segments_entry, values_entry, compression: int): bool throws { - // get the array of segments - var segments = segments_entry.a; - - var prefix: string; - var extension: string; - - (prefix, extension) = getFileMetadata(filename); - - // Generate the filenames based upon the number of targetLocales. - var filenames = generateFilenames(prefix, extension, segments.targetLocales().size); - - //Generate a list of matching filenames to test against. - var matchingFilenames = getMatchingFilenames(prefix, extension); - - var filesExist = processParquetFilenames(filenames, matchingFilenames, TRUNCATE); - - const extraOffset = values_entry.size; - const lastOffset = if segments.size == 0 then 0 else segments[segments.domain.high]; // prevent index error when empty - const lastValIdx = values_entry.a.domain.high; - ref olda = values_entry.a; - - // pull values to the locale of the offset - coforall (loc, idx) in zip(segments.targetLocales(), filenames.domain) with (ref olda) do on loc { - const myFilename = filenames[idx]; - const locDom = segments.localSubdomain(); - var dims: [0..#1] int; - dims[0] = locDom.size: int; - - if (locDom.isEmpty() || locDom.size <= 0) { - // we know append is not supported so creating new empty file - createEmptyListParquetFile(myFilename, dsetName, c_dtype, compression); - } else { - var localSegments = segments[locDom]; - var startValIdx = localSegments[locDom.low]; - - var endValIdx = if (lastOffset == localSegments[locDom.high]) then lastValIdx else segments[locDom.high + 1] - 1; - - var valIdxRange = startValIdx..endValIdx; - writeSegArrayComponent(myFilename, dsetName, olda, valIdxRange, segments, locDom, extraOffset, lastOffset, lastValIdx, c_dtype, compression); - } - } - return filesExist; // trigger warning if overwrite occuring + return writeListColumn(filename, dsetName, segments_entry.a, + values_entry.a, compression: CompressionType); } proc writeStrSegArrayParquet(filename: string, dsetName: string, segments_entry, values_entry, compression: int): bool throws { - extern proc c_writeStrListColumnToParquet(filename, segs_chpl, offsets_chpl, arr_chpl, - dsetname, numelems, rowGroupSize, - dtype, compression, errMsg): int; - // get the array of segments - var segments = segments_entry.a; - - var prefix: string; - var extension: string; - - (prefix, extension) = getFileMetadata(filename); - - // Generate the filenames based upon the number of targetLocales. - var filenames = generateFilenames(prefix, extension, segments.targetLocales().size); - - //Generate a list of matching filenames to test against. - var matchingFilenames = getMatchingFilenames(prefix, extension); - - var filesExist = processParquetFilenames(filenames, matchingFilenames, TRUNCATE); // we know append is not supported - - // Note - seg/segment refers to segarray offsets and off/offsets refers to string object offsets - ref oldOff = values_entry.offsetsEntry.a; - ref oldVal = values_entry.bytesEntry.a; - const extraSegment = values_entry.offsetsEntry.size; - const extraOffset = values_entry.bytesEntry.size; - const lastOffset = if segments.size == 0 then 0 else segments[segments.domain.high]; // prevent index error when empty - const lastOffsetIdx = oldOff.domain.high; - const lastValIdx = oldVal.domain.high; - - // pull values to the locale of the offset - coforall (loc, idx) in zip(segments.targetLocales(), filenames.domain) with (ref oldOff, ref oldVal) do on loc { - const myFilename = filenames[idx]; - - const locDom = segments.localSubdomain(); - - if (locDom.isEmpty() || locDom.size <= 0) { - // we know append is not supported so creating new empty file - var c_dtype = ARROWSTRING; - createEmptyListParquetFile(myFilename, dsetName, c_dtype, compression); - } - else { - var localSegments = segments[locDom]; - var locSegments: [0..#locDom.size+1] int; - locSegments[0..#locDom.size] = segments[locDom]; - if locDom.high == segments.domain.high then - locSegments[locSegments.domain.high] = extraSegment; - else - locSegments[locSegments.domain.high] = segments[locDom.high+1]; - - var startOffsetIdx = localSegments[locDom.low]; - var endOffsetIdx = if (lastOffset == localSegments[locDom.high]) then lastOffsetIdx else segments[locDom.high + 1] - 1; - var offIdxRange = startOffsetIdx..endOffsetIdx; - - var pqErr = new parquetErrorMsg(); - var dtypeRep = ARROWSTRING; - var valPtr: c_ptr(void) = nil; - var offPtr: c_ptr(void) = nil; - - // need to get the local string values - if offIdxRange.size > 0 { - var localOffsets: [offIdxRange] int = oldOff[offIdxRange]; - var startValIdx = oldOff[offIdxRange.low]; - var endValIdx = if (lastOffsetIdx == offIdxRange.high) then lastValIdx else oldOff[offIdxRange.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - var localVals: [valIdxRange] uint(8) = oldVal[valIdxRange]; - - var locOffsets: [0..#offIdxRange.size+1] int; - locOffsets[0..#offIdxRange.size] = oldOff[offIdxRange]; - - if offIdxRange.high == oldOff.domain.high { - locOffsets[locOffsets.domain.high] = extraOffset; - } else { - locOffsets[locOffsets.domain.high] = oldOff[offIdxRange.high+1]; - } - - if localVals.size > 0 { - valPtr = c_ptrTo(localVals); - } - if locOffsets.size > 0 { - offPtr = c_ptrTo(locOffsets); - } - // the call to c must be within the if block so the arrays stay in scope - if c_writeStrListColumnToParquet(myFilename.localize().c_str(), c_ptrTo(locSegments), offPtr, - valPtr, dsetName.localize().c_str(), locSegments.size-1, - ROWGROUPS, dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - else { - // empty segment case - if c_writeStrListColumnToParquet(myFilename.localize().c_str(), c_ptrTo(locSegments), offPtr, - valPtr, dsetName.localize().c_str(), locSegments.size-1, - ROWGROUPS, dtypeRep, compression, c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - } - - } - } - return filesExist; // trigger warning if overwrite occuring + return writeStrListColumn(filename, dsetName, segments_entry.a, + values_entry.offsetsEntry.a, + values_entry.bytesEntry.a, + compression: CompressionType); } proc segarray_toParquetMsg(msgArgs: MessageArgs, st: borrowed SymTab): bool throws { @@ -2031,446 +995,149 @@ module ParquetMsg { } } - proc writeMultiColParquet(filename: string, col_names: [] string, - ncols: int, sym_names: [] string, col_objTypes: [] string, targetLocales: [] locale, - compression: int, st: borrowed SymTab): bool throws { - - extern proc c_writeMultiColToParquet(filename, column_names, ptr_arr, offset_arr, objTypes, - datatypes, segArr_sizes, colnum, numelems, rowGroupSize, compression, errMsg): int; - - var prefix: string; - var extension: string; - (prefix, extension) = getFileMetadata(filename); - - // Generate the filenames based upon the number of targetLocales. - var filenames = generateFilenames(prefix, extension, targetLocales.size); - - //Generate a list of matching filenames to test against. - var matchingFilenames = getMatchingFilenames(prefix, extension); - - // TODO when APPEND is fully deprecated update this to not need the mode. - var filesExist = processParquetFilenames(filenames, matchingFilenames, TRUNCATE); // set to truncate. We will not be supporting appending. - - coforall (loc, idx) in zip(targetLocales, filenames.domain) do on loc { - var pqErr = new parquetErrorMsg(); - const fname = filenames[idx]; - - var ptrList: [0..#ncols] c_ptr(void); - var segmentPtr: [0..#ncols] c_ptr(void); // ptrs to offsets for SegArray. Know number of rows so we know where to stop - var objTypes: [0..#ncols] int; // ObjType enum integer values - var datatypes: [0..#ncols] int; - var sizeList: [0..#ncols] int; - var segarray_sizes: [0..#ncols] int; // track # of values in each column. Used to determine last segment size. - - var my_column_names = col_names; - var c_names: [0..#ncols] c_ptrConst(c_char); - - var segment_ct: [0..#ncols] int; - var seg_sizes_str: [0..#ncols] int; // Track the sizes of string columns and # of segments in segarray str column - var val_sizes_str: [0..#ncols] int; // Track # of values making strings coming to locale - var seg_sizes_int: [0..#ncols] int; // only fill in sizes for int segarray columns - var seg_sizes_uint: [0..#ncols] int; // only fill in sizes for uint segarray columns - var seg_sizes_real: [0..#ncols] int; // only fill in sizes for float segarray columns - var seg_sizes_bool: [0..#ncols] int; // only fill in sizes for bool segarray columns - forall (i, column, ot) in zip(0..#ncols, sym_names, col_objTypes) { - var x: int; - var objType = ot.toUpper(): ObjType; - - if objType == ObjType.STRINGS { - var entry = st[column]; - var e: SegStringSymEntry = toSegStringSymEntry(entry); - var segStr = new SegString("", e); - ref ss = segStr; - var lens = ss.getLengths(); - const locDom = ss.offsets.a.localSubdomain(); - for d in locDom do x += lens[d]; - seg_sizes_str[i] = x; - } - else if objType == ObjType.SEGARRAY { - // parse the json in column to get the component pdarrays - var components: map(string, string) = jsonToMap(column); - var seg_entry = getGenericTypedArrayEntry(components["segments"], st); - var segments = toSymEntry(seg_entry, int); - ref sa = segments.a; - const saD = sa.domain; - var lens: [saD] int; - const high = saD.high; - const locDom = sa.localSubdomain(); - var values = getGenericTypedArrayEntry(components["values"], st); - - segment_ct[i] += locDom.size; - if values.dtype == DType.Strings && locDom.size > 0 { - var e: SegStringSymEntry = toSegStringSymEntry(values); - var segStr = new SegString("", e); - ref ss = segStr; - var lens = ss.getLengths(); - const lastOffset = if sa.size == 0 then 0 else sa[high]; - const lastOffsetIdx = segStr.offsets.a.domain.high; - var startOffsetIdx = sa[locDom.low]; - var endOffsetIdx = if (lastOffset == sa[locDom.high]) then lastOffsetIdx else sa[locDom.high + 1] - 1; - var offIdxRange = startOffsetIdx..endOffsetIdx; - var str_bytes: int; - for d in offIdxRange do str_bytes += lens[d]; - - seg_sizes_str[i] = str_bytes; + private proc registerParquetColumns(ref op, colNames: [] string, + symNames: [] string, + colObjTypes: [] string, + st: borrowed SymTab) throws { + for (colName, symName, objTypeName) in + zip(colNames, symNames, colObjTypes) { + select objTypeName.toUpper(): ObjType { + when ObjType.PDARRAY { + const entry = getGenericTypedArrayEntry(symName, st); + select entry.dtype { + when DType.Int64 do + op.registerColumn(toSymEntry(entry, int).a, colName); + when DType.UInt64 do + op.registerColumn(toSymEntry(entry, uint).a, colName); + when DType.Float64 do + op.registerColumn(toSymEntry(entry, real).a, colName); + when DType.Bool do + op.registerColumn(toSymEntry(entry, bool).a, colName); + otherwise do + throw getErrorWithContext(getL(), getM(), getR(), + msg="Unsupported PDArray DType for writing to Parquet, " + + entry.dtype:string, + errorClass='DataTypeError'); } + } + when ObjType.STRINGS { + const entry = toSegStringSymEntry(st[symName]); + op.registerStrColumn(entry.offsetsEntry.a, entry.bytesEntry.a, + colName); + } + when ObjType.SEGARRAY { + const components = jsonToMap(symName); + const segments = toSymEntry( + getGenericTypedArrayEntry(components["segments"], st), int); + const values = getGenericTypedArrayEntry(components["values"], st); - lens = [(i, s) in zip (saD, sa)] if i == high then values.size - s else sa[i+1] - s; - for d in locDom do x += lens[d]; select values.dtype { - when DType.Int64 { - seg_sizes_int[i] = x; - } - when DType.UInt64 { - seg_sizes_uint[i] = x; - } - when DType.Float64 { - seg_sizes_real[i] = x; - } - when DType.Bool { - seg_sizes_bool[i] = x; - } + when DType.Int64 do + op.registerListColumn(segments.a, + toSymEntry(values, int).a, colName); + when DType.UInt64 do + op.registerListColumn(segments.a, + toSymEntry(values, uint).a, colName); + when DType.Float64 do + op.registerListColumn(segments.a, + toSymEntry(values, real).a, colName); + when DType.Bool do + op.registerListColumn(segments.a, + toSymEntry(values, bool).a, colName); when DType.Strings { - val_sizes_str[i] = x; + const strings = toSegStringSymEntry(values); + op.registerStrListColumn(segments.a, strings.offsetsEntry.a, + strings.bytesEntry.a, colName); } - otherwise { + otherwise do throw getErrorWithContext(getL(), getM(), getR(), - msg="Unsupported SegArray DType for writing to Parquet, ".format(values.dtype: string), - errorClass='DataTypeError'); - } - } - } - } - - var totalSegs: int = + reduce segment_ct; // total # of offsets on locale - var segment_tracking: [0..#totalSegs] int; // array to write offset values into after adjusting for locale - var segment_idx = (+ scan segment_ct) - segment_ct; // offset start indexes for each column - - var locSize_str: int = + reduce seg_sizes_str; - var str_vals: [0..#locSize_str] uint(8); - var locSize_int: int = + reduce seg_sizes_int; - var int_vals: [0..#locSize_int] int; - var locSize_uint: int = + reduce seg_sizes_uint; - var uint_vals: [0..#locSize_uint] uint; - var locSize_real: int = + reduce seg_sizes_real; - var real_vals: [0..#locSize_real] real; - var locSize_bool: int = + reduce seg_sizes_bool; - var bool_vals: [0..#locSize_bool] bool; - - // indexes for which values go to which columns - var str_idx = (+ scan seg_sizes_str) - seg_sizes_str; - var int_idx = (+ scan seg_sizes_int) - seg_sizes_int; - var uint_idx = (+ scan seg_sizes_uint) - seg_sizes_uint; - var real_idx = (+ scan seg_sizes_real) - seg_sizes_real; - var bool_idx = (+ scan seg_sizes_bool) - seg_sizes_bool; - - // populate data based on object and data types - forall (i, column, ot, si, ii, ui, ri, bi, segidx) in zip(0..#ncols, sym_names, col_objTypes, str_idx, int_idx, uint_idx, real_idx, bool_idx, segment_idx) { - // generate the local c string list of column names - c_names[i] = my_column_names[i].localize().c_str(); - - select ot.toUpper(): ObjType { - when ObjType.STRINGS { - var entry = st[column]; - var e: SegStringSymEntry = toSegStringSymEntry(entry); - var segStr = new SegString("", e); - ref ss = segStr; - var A = ss.offsets.a; - const lastOffset = if A.size == 0 then 0 else A[A.domain.high]; // prevent index error when empty - const lastValIdx = ss.values.a.domain.high; - const locDom = ss.offsets.a.localSubdomain(); - - objTypes[i] = ObjType.STRINGS: int; - datatypes[i] = ARROWSTRING; - - if locDom.size > 0 { - var localOffsets = A[locDom]; - var startValIdx = localOffsets[locDom.low]; - var endValIdx = if (lastOffset == localOffsets[locDom.high]) then lastValIdx else A[locDom.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - ref olda = ss.values.a; - str_vals[si..#valIdxRange.size] = olda[valIdxRange]; - ptrList[i] = c_ptrTo(str_vals[si]): c_ptr(void); - sizeList[i] = locDom.size; - } - } - when ObjType.SEGARRAY { - // parse the json in column to get the component pdarrays - var components: map(string, string) = jsonToMap(column); - - // access segments symentry - var seg_entry = getGenericTypedArrayEntry(components["segments"], st); - var segments = toSymEntry(seg_entry, int); - - ref S = segments.a; - const locDom = segments.a.localSubdomain(); - objTypes[i] = ObjType.SEGARRAY: int; - - if locDom.size > 0 { - const lastSegment = if S.size == 0 then 0 else S[S.domain.high]; // prevent index error when empty; - const localSegments = S[locDom]; - const startValIdx = localSegments[locDom.low]; - sizeList[i] = locDom.size; - segment_tracking[segidx..#locDom.size] = localSegments - startValIdx; - segmentPtr[i] = c_ptrTo(segment_tracking[segidx]); - - var valEntry = getGenericTypedArrayEntry(components["values"], st); - select valEntry.dtype { - when DType.Int64 { - segarray_sizes[i] = seg_sizes_int[i]; - var values = toSymEntry(valEntry, int); - const lastValIdx = values.a.domain.high; - - datatypes[i] = ARROWINT64; - - const endValIdx = if (lastSegment == localSegments[locDom.high]) then lastValIdx else S[locDom.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - ref olda = values.a; - if !int_vals.domain.isEmpty() { - int_vals[ii..#valIdxRange.size] = olda[valIdxRange]; - ptrList[i] = c_ptrTo(int_vals[ii]): c_ptr(void); - } - } - when DType.UInt64 { - segarray_sizes[i] = seg_sizes_uint[i]; - var values = toSymEntry(valEntry, uint); - const lastValIdx = values.a.domain.high; - - datatypes[i] = ARROWUINT64; - - var endValIdx = if (lastSegment == localSegments[locDom.high]) then lastValIdx else S[locDom.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - ref olda = values.a; - if !uint_vals.domain.isEmpty() { - uint_vals[ui..#valIdxRange.size] = olda[valIdxRange]; - ptrList[i] = c_ptrTo(uint_vals[ui]): c_ptr(void); - } - } - when DType.Float64 { - segarray_sizes[i] = seg_sizes_real[i]; - var values = toSymEntry(valEntry, real); - const lastValIdx = values.a.domain.high; - - datatypes[i] = ARROWDOUBLE; - - var endValIdx = if (lastSegment == localSegments[locDom.high]) then lastValIdx else S[locDom.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - ref olda = values.a; - if !real_vals.domain.isEmpty() { - real_vals[ri..#valIdxRange.size] = olda[valIdxRange]; - ptrList[i] = c_ptrTo(real_vals[ri]): c_ptr(void); - } - } - when DType.Bool { - segarray_sizes[i] = seg_sizes_bool[i]; - var values = toSymEntry(valEntry, bool); - const lastValIdx = values.a.domain.high; - - datatypes[i] = ARROWBOOLEAN; - - var endValIdx = if (lastSegment == localSegments[locDom.high]) then lastValIdx else S[locDom.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - ref olda = values.a; - if !bool_vals.domain.isEmpty() { - bool_vals[bi..#valIdxRange.size] = olda[valIdxRange]; - ptrList[i] = c_ptrTo(bool_vals[bi]): c_ptr(void); - } - } - when DType.Strings { - segarray_sizes[i] = val_sizes_str[i]; - var values = toSegStringSymEntry(valEntry); - ref oldOff = values.offsetsEntry.a; - ref oldVal = values.bytesEntry.a; - - const extraSegment = oldOff.size; - const extraOffset = oldVal.size; - const lastSegment = if S.size == 0 then 0 else S[S.domain.high]; - const lastOffsetIdx = oldOff.domain.high; - const lastValIdx = oldVal.domain.high; - - datatypes[i] = ARROWSTRING; - - var startOffsetIdx = localSegments[locDom.low]; - var endOffsetIdx = if (lastSegment == localSegments[locDom.high]) then lastOffsetIdx else S[locDom.high + 1] - 1; - var offIdxRange = startOffsetIdx..endOffsetIdx; - if offIdxRange.size > 0 { - - var localOffsets: [offIdxRange] int = oldOff[offIdxRange]; - var startValIdx = localOffsets[offIdxRange.low]; - - var endValIdx = if (lastOffsetIdx == offIdxRange.high) then lastValIdx else oldOff[offIdxRange.high + 1] - 1; - var valIdxRange = startValIdx..endValIdx; - if !str_vals.domain.isEmpty() { - str_vals[si..#valIdxRange.size] = oldVal[valIdxRange]; - ptrList[i] = c_ptrTo(str_vals[si]): c_ptr(void); - } - } - } - otherwise { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Unsupported SegArray DType for writing to Parquet, ".format(valEntry.dtype: string), - errorClass='DataTypeError'); - } - } - } - else { - // set the datatype for empty locales to ensure that metadata is correct in all files - var valEntry = getGenericTypedArrayEntry(components["values"], st); - select valEntry.dtype { - when DType.Int64 { - datatypes[i] = ARROWINT64; - } - when DType.UInt64 { - datatypes[i] = ARROWUINT64; - } - when DType.Float64 { - datatypes[i] = ARROWDOUBLE; - } - when DType.Bool { - datatypes[i] = ARROWBOOLEAN; - } - when DType.Strings { - datatypes[i] = ARROWSTRING; - } - otherwise { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Unsupported SegArray DType for writing to Parquet, ".format(valEntry.dtype: string), - errorClass='DataTypeError'); - } - } - } - } - when ObjType.PDARRAY { - var entry = getGenericTypedArrayEntry(column, st); - select entry.dtype { - when DType.Int64 { - var e = toSymEntry(toGenSymEntry(entry), int); - var locDom = e.a.localSubdomain(); - objTypes[i] = ObjType.PDARRAY: int; - datatypes[i] = ARROWINT64; - // set the pointer to the entry array in the list of Pointers - if locDom.size > 0 { - ptrList[i] = c_ptrTo(e.a[locDom.low]): c_ptr(void); - sizeList[i] = locDom.size; - } - } - when DType.UInt64 { - var e = toSymEntry(toGenSymEntry(entry), uint); - var locDom = e.a.localSubdomain(); - objTypes[i] = ObjType.PDARRAY: int; - datatypes[i] = ARROWUINT64; - // set the pointer to the entry array in the list of Pointers - if locDom.size > 0 { - ptrList[i] = c_ptrTo(e.a[locDom.low]): c_ptr(void); - sizeList[i] = locDom.size; - } - } - when DType.Float64 { - var e = toSymEntry(toGenSymEntry(entry), real); - var locDom = e.a.localSubdomain(); - objTypes[i] = ObjType.PDARRAY: int; - datatypes[i] = ARROWDOUBLE; - // set the pointer to the entry array in the list of Pointers - if locDom.size > 0 { - ptrList[i] = c_ptrTo(e.a[locDom.low]): c_ptr(void); - sizeList[i] = locDom.size; - } - } - when DType.Bool { - var e = toSymEntry(toGenSymEntry(entry), bool); - var locDom = e.a.localSubdomain(); - objTypes[i] = ObjType.PDARRAY: int; - datatypes[i] = ARROWBOOLEAN; - // set the pointer to the entry array in the list of Pointers - if locDom.size > 0 { - ptrList[i] = c_ptrTo(e.a[locDom.low]): c_ptr(void); - sizeList[i] = locDom.size; - } - } - otherwise { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Unsupported PDArray DType for writing to Parquet, ".format(entry.dtype: string), + msg="Unsupported SegArray DType for writing to Parquet, " + + values.dtype:string, errorClass='DataTypeError'); - } - } - } - otherwise { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Writing Parquet files (multi-column) does not support %s columns.".format(ot), - errorClass='DataTypeError' - ); } } - } - // validate all elements same size - var numelems: int = sizeList[0]; - if !(&& reduce (sizeList==numelems)) { - throw getErrorWithContext(getL(), getM(), getR(), - msg="Parquet columns must be the same size", - errorClass='WriteModeError' - ); - } - - var result: int = c_writeMultiColToParquet(fname.localize().c_str(), c_ptrTo(c_names), c_ptrTo(ptrList), c_ptrTo(segmentPtr), c_ptrTo(objTypes), c_ptrTo(datatypes), c_ptrTo(segarray_sizes), ncols, numelems, ROWGROUPS, compression, c_ptrTo(pqErr.errMsg)); - if result == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); + otherwise do + throw getErrorWithContext(getL(), getM(), getR(), + msg="Writing Parquet files does not support " + + objTypeName + " columns.", + errorClass='DataTypeError'); } } + } + + private proc writeMultiColWithOps(filename: string, + colNames: [] string, + symNames: [] string, + colObjTypes: [] string, + compression: int, + st: borrowed SymTab, + const ref sharedDom) throws { + const filesExist = filesExistForWrite( + filename, sharedDom.targetLocales().size, TRUNCATE); + + var op = new pqWriteOp(filename, sharedDom); + op.compression = compression; + op.distributed = true; + registerParquetColumns(op, colNames, symNames, colObjTypes, st); + op.write(); return filesExist; } - proc identifyTargetLocales(name: string, objType: string, st: borrowed SymTab) throws { - var targetLocales; - select objType.toUpper(): ObjType { + private proc writeMultiColParquet(filename: string, + colNames: [] string, + symNames: [] string, + colObjTypes: [] string, + compression: int, + st: borrowed SymTab) throws { + const firstName = symNames[symNames.domain.low]; + select colObjTypes[colObjTypes.domain.low].toUpper(): ObjType { when ObjType.STRINGS { - var entry = st[name]; - var e: SegStringSymEntry = toSegStringSymEntry(entry); - var segStr = new SegString("", e); - targetLocales = segStr.offsets.a.targetLocales(); + const first = toSegStringSymEntry(st[firstName]); + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + first.offsetsEntry.a.domain); } when ObjType.SEGARRAY { - // parse the json in column to get the component pdarrays - var components: map(string, string) = jsonToMap(name); - - // access segments symentry - var seg_entry = getGenericTypedArrayEntry(components["segments"], st); - var segments = toSymEntry(seg_entry, int); - targetLocales = segments.a.targetLocales(); + const components = jsonToMap(firstName); + const first = toSymEntry( + getGenericTypedArrayEntry(components["segments"], st), int); + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + first.a.domain); } when ObjType.PDARRAY { - var entry = st[name]; - var entryDtype = (entry: borrowed GenSymEntry).dtype; - select entryDtype { - when DType.Int64 { - var e = toSymEntry(toGenSymEntry(entry), int); - targetLocales = e.a.targetLocales(); - } - when DType.UInt64 { - var e = toSymEntry(toGenSymEntry(entry), uint); - targetLocales = e.a.targetLocales(); - } - when DType.Float64 { - var e = toSymEntry(toGenSymEntry(entry), real); - targetLocales = e.a.targetLocales(); - } - when DType.Bool { - var e = toSymEntry(toGenSymEntry(entry), bool); - targetLocales = e.a.targetLocales(); - } - otherwise { + const first = getGenericTypedArrayEntry(firstName, st); + select first.dtype { + when DType.Int64 do + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + toSymEntry(first, int).a.domain); + when DType.UInt64 do + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + toSymEntry(first, uint).a.domain); + when DType.Float64 do + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + toSymEntry(first, real).a.domain); + when DType.Bool do + return writeMultiColWithOps(filename, colNames, symNames, + colObjTypes, compression, st, + toSymEntry(first, bool).a.domain); + otherwise do throw getErrorWithContext(getL(), getM(), getR(), - msg="Writing Parquet files (multi-column) does not support columns of type %s".format(entryDtype: string), - errorClass='DataTypeError'); - } + msg="Unsupported PDArray DType for writing to Parquet, " + + first.dtype:string, + errorClass='DataTypeError'); } } - otherwise { + otherwise do throw getErrorWithContext(getL(), getM(), getR(), - msg="Writing Parquet files (multi-column) does not support %s columns.".format(objType), - errorClass='DataTypeError'); - } + msg="Writing Parquet files does not support " + + colObjTypes[colObjTypes.domain.low] + " columns.", + errorClass='DataTypeError'); } - return targetLocales; + return false; } proc toParquetMultiColMsg(cmd: string, msgArgs: borrowed MessageArgs, st: borrowed SymTab): MsgTuple throws { @@ -2489,12 +1156,10 @@ module ParquetMsg { // compression format as integer const compression = msgArgs.getValueOf("compression").toUpper(): CompressionType; - // use the first entry to identify target locales. Assuming all have same distribution - var targetLocales = identifyTargetLocales(sym_names[0], col_objType_strs[0], st); - var warnFlag: bool; try { - warnFlag = writeMultiColParquet(filename, col_names, ncols, sym_names, col_objType_strs, targetLocales, compression:int, st); + warnFlag = writeMultiColParquet(filename, col_names, sym_names, + col_objType_strs, compression:int, st); } catch e: FileNotFoundError { var errorMsg = "Unable to open %s for writing: %s".format(filename,e.message()); pqLogger.error(getM(),getR(),getL(),errorMsg); @@ -2555,23 +1220,7 @@ module ParquetMsg { } try { - // TODO use getDatasets - extern proc c_getDatasetNames(filename, dsetResult, readNested, errMsg): int(32); - extern proc strlen(a): int; - var pqErr = new parquetErrorMsg(); - var res: c_ptr(uint(8)); - defer { - extern proc c_free_string(ptr); - c_free_string(res); - } - if c_getDatasetNames(filename.c_str(), c_ptrTo(res), read_nested, - c_ptrTo(pqErr.errMsg)) == ARROWERROR { - pqErr.parquetError(getL(), getR(), getM()); - } - try! repMsg = string.createCopyingBuffer(res, strlen(res)); - var items = new list(repMsg.split(",")); // convert to json - - repMsg = formatJson(items); + repMsg = formatJson(getDatasets(filename, readNested=read_nested)); } catch e : Error { var errorMsg = "Failed to process Parquet file %?".format(e.message()); return new MsgTuple(errorMsg, MsgType.ERROR); @@ -2697,5 +1346,11 @@ module ParquetMsg { registerFunction("writeParquet", toparquetMsg, getM()); registerFunction("lspq", lspqMsg, getM()); registerFunction("getnullparquet", nullIndicesMsg, getM()); - ServerConfig.appendToConfigStr("ARROW_VERSION", getVersionInfo()); + try { + ServerConfig.appendToConfigStr("ARROW_VERSION", getVersionInfo()); + } catch e: Error { + ServerConfig.appendNullToConfigStr("ARROW_VERSION"); + pqLogger.error(getM(), getR(), getL(), + "Failed to retrieve Arrow version: ", e.message()); + } } diff --git a/src/ServerConfig.chpl b/src/ServerConfig.chpl index e97424db04c..bc955dbfa4e 100644 --- a/src/ServerConfig.chpl +++ b/src/ServerConfig.chpl @@ -486,6 +486,12 @@ module ServerConfig cfgStr = tmp_json + "," + Q + key + QCQ + val + Q + "}"; } + proc appendNullToConfigStr(key:string) { + var idx_close = cfgStr.rfind("}"):int; + var tmp_json = cfgStr(0..idx_close-1); + cfgStr = tmp_json + "," + Q + key + Q + ":null}"; + } + /* proposed replacement for `timeSinceEpoch().totalSeconds()` */ proc currentTime() { use Time; diff --git a/src/parquet/ReadParquet.cpp b/src/parquet/ReadParquet.cpp deleted file mode 100644 index af1b53cb11f..00000000000 --- a/src/parquet/ReadParquet.cpp +++ /dev/null @@ -1,1225 +0,0 @@ -#include "ReadParquet.h" - - -namespace akcpp { - -namespace { - -// ChplBuf::ElemType can be used to get the type of the Chapel buffer from an -// Arrow type. -template -struct ChplBuf { - using ElemType = typename ArrowType::c_type; -}; - -// specialization for storing 32-bit data in 64-bit arrays -template<> struct ChplBuf { using ElemType = int64_t; }; -template<> struct ChplBuf { using ElemType = uint64_t; }; -template<> struct ChplBuf { using ElemType = double; }; - -// TODO The implementation seemingly reads 128-bit decimal values into -// Chapel `real`s (or C `double`s). I don't see how this can work properly with -// large enough data. For now, I think this keeps the implementation practically -// same as the old, likely carrying the same bug into the new implementation, -// unfortunately. -// See https://github.com/Bears-R-Us/arkouda/issues/4911. -template<> struct ChplBuf {using ElemType = double;}; - -// this struct is used to bundle different types into one, based on a single -// ArrowType -template -struct TypeBundle { - using ReaderType = typename parquet::TypedColumnReader; - using ChplType = typename ChplBuf::ElemType; - using PqType = typename ArrowType::c_type; -}; - -template<> -struct TypeBundle { - using ReaderType = typename parquet::TypedColumnReader; - using ChplType = typename ChplBuf::ElemType; - using PqType = typename parquet::FixedLenByteArray; -}; -} // end anonymous namespace - -// Generic read implementation -template -int64_t ColReadOp::read() { - using Types = TypeBundle; - using ChplType = typename Types::ChplType; - using ReaderType = typename Types::ReaderType; - - auto chpl_ptr = (ChplType*)chpl_arr; - int64_t num_read = 0; - ReaderType* reader = static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - // TODO find better variable names. values_read and num_read are confusing - int64_t values_read = 0; - - // note that floats are the types that natively support nulls. This generic - // read() shouldn't do anything with floats - if (nullMode == noNulls || nullMode == onlyFloats) { - while (reader->HasNext() && row_idx < numElems) { - if((numElems - row_idx) < batchSize) // adjust batchSize if needed - batchSize = numElems - row_idx; - std::ignore = reader->ReadBatch(batchSize, nullptr, nullptr, - &chpl_ptr[row_idx], &values_read); - row_idx+=values_read; - num_read+=values_read; - } - } - else { - int16_t definition_level; // nullable type and only reading single records in batch - while (reader->HasNext() && row_idx < numElems) { - std::ignore = reader->ReadBatch(1, &definition_level, nullptr, - &chpl_ptr[row_idx], &values_read); - // if values_read is 0, that means that it was a null value - if(values_read == 0) { - where_null_chpl[row_idx] = true; - } - row_idx++; - num_read++; - } - } - return num_read; -} - -template -int64_t ColReadOp::_readShortIntegral() { - using ReaderType = typename Types::ReaderType; - using ChplType = typename Types::ChplType; - using PqType = typename Types::PqType; - - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - // TODO find better variable names. values_read and num_read are confusing - int64_t values_read = 0; - - int64_t num_read = 0; - if (not hasNonFloatNulls) { - // TODO we can read into the actual Chapel array, and then space the data - // out in reverse order. I don't think we need to have this temporary - // buffer - PqType* tmpArr = (PqType*)malloc(batchSize * sizeof(PqType)); - while (reader->HasNext() && row_idx < numElems) { - if((numElems - row_idx) < batchSize) // adjust batchSize if needed - batchSize = numElems - row_idx; - - // Can't read directly into chpl_ptr because it is int64 - std::ignore = reader->ReadBatch(batchSize, nullptr, nullptr, tmpArr, - &values_read); - - for (int64_t j = 0; j < values_read; j++) { - chpl_ptr[row_idx+j] = (ChplType)tmpArr[j]; - } - - row_idx+=values_read; - num_read+=values_read; - } - free(tmpArr); - } - else { - PqType tmp; - // Engin: we don't seem to use this anywhere, but passing nullptr instead - // of this causes errors. Why? - int16_t definition_level; // nullable type and only reading single - // records in batch - while (reader->HasNext() && row_idx < numElems) { - std::ignore = reader->ReadBatch(1, &definition_level, nullptr, &tmp, - &values_read); - // if values_read is 0, that means that it was a null value - if(values_read == 0) { - where_null_chpl[row_idx] = true; - } - else { - chpl_ptr[row_idx] = (ChplType)tmp; - } - row_idx++; - num_read++; - } - } - return num_read; -} - -template<> -int64_t ColReadOp::read() { - return _readShortIntegral>(); -} - -template<> -int64_t ColReadOp::read() { - return _readShortIntegral>(); -} - -template<> -int64_t ColReadOp::read() { - using Types = TypeBundle; - using ReaderType = typename Types::ReaderType; - using ChplType = typename Types::ChplType; - using PqType = typename Types::PqType; - - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - // TODO find better variable names. values_read and num_read are confusing - int64_t values_read = 0; - - int64_t num_read = 0; - if (nullMode == noNulls) { - PqType* tmpArr = (PqType*)malloc(batchSize * sizeof(PqType)); - while (reader->HasNext() && row_idx < numElems) { - if((numElems - row_idx) < batchSize) // adjust batchSize if needed - batchSize = numElems - row_idx; - std::ignore = reader->ReadBatch(batchSize, nullptr, nullptr, tmpArr, - &values_read); - - // promote to larger type - for (int64_t j = 0; j < values_read; j++) { - chpl_ptr[row_idx+j] = (ChplType)tmpArr[j]; - } - - row_idx+=values_read; - num_read+=values_read; - - } - free(tmpArr); - } - else { - int16_t definition_level; // nullable type and only reading single records in batch - while (reader->HasNext() && row_idx < numElems) { - PqType value = 0; - std::ignore = 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) { - chpl_ptr[row_idx] = (ChplType) value; - } - else { - chpl_ptr[row_idx] = NAN; - } - row_idx++; - num_read++; - } - } - return num_read; -} - -template<> -int64_t ColReadOp::read() { - using Types = TypeBundle; - using ReaderType = typename Types::ReaderType; - using ChplType = typename Types::ChplType; - using PqType = typename Types::PqType; - - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - // TODO find better variable names. values_read and num_read are confusing - int64_t values_read = 0; - - int64_t num_read = 0; - if (nullMode == noNulls) { - while (reader->HasNext() && row_idx < numElems) { - if((numElems - row_idx) < batchSize) // adjust batchSize if needed - batchSize = numElems - row_idx; - std::ignore = reader->ReadBatch(batchSize, nullptr, nullptr, - &chpl_ptr[row_idx], &values_read); - row_idx+=values_read; - num_read+=values_read; - } - } - else { - int16_t definition_level; // nullable type and only reading single records in batch - while (reader->HasNext() && row_idx < numElems) { - PqType value = 0; - std::ignore = 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) { - chpl_ptr[row_idx] = (ChplType) value; - } - else { - chpl_ptr[row_idx] = NAN; - } - row_idx++; - num_read++; - } - } - return num_read; -} - -// For Parquet's "Decimal" types -- typically larger, or more precise floats -template<> -int64_t ColReadOp::read() { - using Types = TypeBundle; - using ReaderType = typename Types::ReaderType; - using ChplType = typename Types::ChplType; - using PqType = typename Types::PqType; - - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = static_cast(column_reader.get()); - startIdx -= reader->Skip(*startIdx); - - // TODO find better variable names. values_read and num_read are confusing - int64_t values_read = 0; - int64_t num_read = 0; - - using LogType = parquet::DecimalLogicalType; - const auto& type = dynamic_cast(*col_info->logical_type()); - const int64_t precision = type.precision(); - - // In ReadParquet.cpp, there is a basic look up table for this. But number - // of required bytes can be found mathematically: - int numbits = ceil(precision*3.321928); // the magic number is the constant - // from log2(10^precision) - if (numbits%8==0) numbits++; // add a bit for the sign bit if we - // are at the byte boundary - const auto numbytes = ceil(numbits/8.0); - - while (reader->HasNext() && row_idx < numElems) { - PqType value; - std::ignore = reader->ReadBatch(1, nullptr, nullptr, &value, - &values_read); - arrow::Decimal128 v; - PARQUET_ASSIGN_OR_THROW(v, - ::arrow::Decimal128::FromBigEndian(value.ptr, - numbytes)); - - chpl_ptr[row_idx] = v.ToDouble(0); - row_idx+=values_read; - num_read+=values_read; - } - - return num_read; -} - - -int readAllCols(const char* filename, void** chpl_arrs, int* types, - bool* where_null_chpl, int64_t numElems, int64_t startIdx, - int64_t batchSize, - chplEnum_t nullMode, char** errMsg) { - try { - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = - parquet_reader->metadata(); - - int num_row_groups = file_metadata->num_row_groups(); - - const auto num_cols = cpp_getNumCols(filename, errMsg); - if (num_cols == ARROWERROR) { - return ARROWERROR; - } - - std::vector startIdxPerCol; - startIdxPerCol.resize(num_cols, startIdx); - - int64_t row_idx = 0; - - for (int rg_idx = 0; (rg_idx row_group_reader = - parquet_reader->RowGroup(rg_idx); - - std::shared_ptr column_reader; - - int64_t nrows_in_iter = -1; - - for (int col_idx = 0; col_idxColumn(col_idx); - const parquet::ColumnDescriptor* col_info = - file_metadata->schema()->Column(col_idx); - - void* chpl_arr = chpl_arrs[col_idx]; - - // TODO, I want values of this type to be created per read operation, - // not per column, per rowgroup - auto op = ColReadOp { chpl_arr, - &startIdxPerCol[col_idx], - column_reader, - false, // has_non_float_nulls is for backward - // compat and ignored while reading all - // columns. nullMode is used instead. - nullMode, - row_idx, - numElems, - batchSize, - where_null_chpl, - col_info }; - - int64_t nread; - switch (types[col_idx]) { - case ARROWFLOAT: nread = op.read(); break; - case ARROWDOUBLE: nread = op.read(); break; - case ARROWINT64: nread = op.read(); break; - case ARROWUINT64: nread = op.read(); break; - case ARROWBOOLEAN: nread = op.read(); break; - case ARROWINT32: nread = op.read(); break; - case ARROWUINT32: nread = op.read(); break; - case ARROWDECIMAL: nread = op.read(); break; - default: - // TODO we might want to have our own exception types on C++ side, - // too - throw std::domain_error("Unknown Arrow type"); - } - - if (nrows_in_iter == -1) { - // this is the first time we are reading in this row group. We'll need - // to bump up the row_idx in the end of the iteration by nrows_read - nrows_in_iter = nread; - } - else { - // we already now how many rows we are reading per column, which is - // stored in nrows_in_iter. But did we read the correct number of rows - // in this particular iteration? - if (nread != nrows_in_iter) { - std::stringstream msgStream; - msgStream << "Uneven number of rows are read."; - msgStream << " Expected " << nrows_in_iter; - msgStream << " But read " << nread << " instead\n"; - throw std::length_error(msgStream.str()); - return ARROWERROR; - } - } - } - - row_idx += nrows_in_iter; - } - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} -} // end namespace akcpp - -// Returns the number of elements read -template -int64_t readColumn(void* chpl_arr, int64_t *startIdx, std::shared_ptr column_reader, - bool hasNonFloatNulls, int64_t i, int64_t numElems, int64_t batchSize, - int64_t values_read, bool* where_null_chpl) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (ChplType*)chpl_arr; - int64_t num_read = 0; - ReaderType* reader = - static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - if (not hasNonFloatNulls) { - while (reader->HasNext() && i < numElems) { - if((numElems - i) < batchSize) // adjust batchSize if needed - batchSize = numElems - i; - (void)reader->ReadBatch(batchSize, nullptr, nullptr, &chpl_ptr[i], &values_read); - i+=values_read; - num_read += values_read; - } - } - else { - while (reader->HasNext() && i < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &chpl_ptr[i], &values_read); - // if values_read is 0, that means that it was a null value - if(values_read == 0) { - where_null_chpl[i] = true; - } - i++; - num_read++; - } - } - return num_read; -} - -template -int64_t readColumnDbFl(void* chpl_arr, int64_t *startIdx, - std::shared_ptr column_reader, - bool hasNonFloatNulls, int64_t i, int64_t numElems, int64_t batchSize, - int64_t values_read, bool* where_null_chpl) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = - static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - int64_t num_read = 0; - while (reader->HasNext() && i < numElems) { - PqType 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) { - // this means it wasn't null - chpl_ptr[i] = (ChplType) value; - } else { - chpl_ptr[i] = NAN; - } - i++; - num_read++; - } - return num_read; -} - -template -int64_t readColumnIrregularBitWidth(void* chpl_arr, int64_t *startIdx, std::shared_ptr column_reader, - bool hasNonFloatNulls, int64_t i, int64_t numElems, int64_t batchSize, - int64_t values_read, bool* where_null_chpl) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (ChplType*)chpl_arr; - ReaderType* reader = - static_cast(column_reader.get()); - *startIdx -= reader->Skip(*startIdx); - - int64_t num_read = 0; - if (not hasNonFloatNulls) { - PqType* tmpArr = (PqType*)malloc(batchSize * sizeof(PqType)); - while (reader->HasNext() && i < numElems) { - if((numElems - i) < batchSize) // adjust batchSize if needed - batchSize = numElems - i; - - // Can't read directly into chpl_ptr because it is int64 - (void)reader->ReadBatch(batchSize, nullptr, nullptr, - (int32_t*)tmpArr, &values_read); - for (int64_t j = 0; j < values_read; j++) - chpl_ptr[i+j] = (ChplType)tmpArr[j]; - i+=values_read; - num_read+=values_read; - } - free(tmpArr); - } - else { - PqType tmp; - while (reader->HasNext() && i < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, - (int32_t*)&tmp, &values_read); - // if values_read is 0, that means that it was a null value - if(values_read == 0) { - where_null_chpl[i] = true; - } - else { - chpl_ptr[i] = (int64_t)tmp; - } - i++; - num_read++; - } - } - return num_read; -} - -int cpp_readStrColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t batchSize, char** errMsg) { - try { - int64_t ty = cpp_getType(filename, colname, errMsg); - - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - int64_t i = 0; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - auto idx = file_metadata -> schema() -> ColumnIndex(colname); - auto max_def = file_metadata -> schema() -> Column(idx) -> max_definition_level(); // needed to determine if nulls are allowed - - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - - column_reader = row_group_reader->Column(idx); - - if(ty == ARROWSTRING) { - auto chpl_ptr = (unsigned char*)chpl_arr; - parquet::ByteArrayReader* reader = - static_cast(column_reader.get()); - - int totalProcessed = 0; - std::vector values(batchSize); - while (reader->HasNext() && totalProcessed < numElems) { - std::vector definition_levels(batchSize,-1); - if((numElems - totalProcessed) < batchSize) // adjust batchSize if needed - batchSize = numElems - totalProcessed; - - (void)reader->ReadBatch(batchSize, definition_levels.data(), nullptr, values.data(), &values_read); - totalProcessed += values_read; - int j = 0; - int numProcessed = 0; - while(j < batchSize) { - if(definition_levels[j] == 1) { - for(int k = 0; k < values[numProcessed].len; k++) { - chpl_ptr[i] = values[numProcessed].ptr[k]; - i++; - } - i++; // skip one space so the strings are null terminated with a 0 - numProcessed++; - } else if(definition_levels[j] == 0) { - i++; - } else { - j = batchSize; // exit loop, not read - } - j++; - } - } - } - } - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - - -int cpp_readColumnByName(const char* filename, void* chpl_arr, bool* where_null_chpl, const char* colname, int64_t numElems, int64_t startIdx, int64_t batchSize, int64_t byteLength, bool hasNonFloatNulls, char** errMsg) { - try { - int64_t ty = cpp_getType(filename, colname, errMsg); - - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - int64_t i = 0; - for (int r = 0; (r < num_row_groups) && (i < numElems); r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - auto idx = file_metadata -> schema() -> ColumnIndex(colname); - - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - auto max_def = file_metadata -> schema() -> Column(idx) -> max_definition_level(); // needed to determine if nulls are allowed - - column_reader = row_group_reader->Column(idx); - - // Since int64 and uint64 Arrow dtypes share a physical type and only differ - // in logical type, they must be read from the file in the same way - if(ty == ARROWINT64 || ty == ARROWUINT64) { - i += readColumn(chpl_arr, - &startIdx, - column_reader, - hasNonFloatNulls, - i, - numElems, - batchSize, - values_read, - where_null_chpl); - } else if(ty == ARROWINT32) { - i += - readColumnIrregularBitWidth( - chpl_arr, &startIdx, column_reader, hasNonFloatNulls, i, - numElems, batchSize, values_read, where_null_chpl); - } else if(ty == ARROWUINT32) { - i += - readColumnIrregularBitWidth( - chpl_arr, &startIdx, column_reader, hasNonFloatNulls, i, - numElems, batchSize, values_read, where_null_chpl); - } else if(ty == ARROWBOOLEAN) { - i += readColumn(chpl_arr, &startIdx, column_reader, hasNonFloatNulls, i, - numElems, batchSize, values_read, where_null_chpl); - } else if(ty == ARROWSTRING) { - 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(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++; - } - } - i++; // skip one space so the strings are null terminated with a 0 - } - } else if(ty == ARROWFLOAT) { - i += readColumnDbFl(chpl_arr, &startIdx, column_reader, hasNonFloatNulls, i, - numElems, batchSize, values_read, where_null_chpl); - } else if(ty == ARROWDOUBLE) { - i += readColumnDbFl(chpl_arr, &startIdx, column_reader, hasNonFloatNulls, i, - numElems, batchSize, values_read, where_null_chpl); - } else if(ty == ARROWDECIMAL) { - auto chpl_ptr = (double*)chpl_arr; - parquet::FixedLenByteArray value; - parquet::FixedLenByteArrayReader* reader = - static_cast(column_reader.get()); - startIdx -= reader->Skip(startIdx); - - while (reader->HasNext() && i < numElems) { - (void)reader->ReadBatch(1, nullptr, nullptr, &value, &values_read); - arrow::Decimal128 v; - PARQUET_ASSIGN_OR_THROW(v, - ::arrow::Decimal128::FromBigEndian(value.ptr, byteLength)); - - chpl_ptr[i] = v.ToDouble(0); - i+=values_read; - } - } - } - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_readListColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg) { - try { - int64_t ty = cpp_getType(filename, colname, errMsg); - if (ty == ARROWLIST){ - int64_t lty = cpp_getListType(filename, colname, errMsg); - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - auto idx = file_metadata -> schema() -> group_node() -> FieldIndex(colname); - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - - int64_t i = 0; - int64_t arrayIdx = 0; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - int16_t definition_level; // needed for any type that is nullable - - std::shared_ptr column_reader = row_group_reader->Column(idx); - if(lty == ARROWINT64 || lty == ARROWUINT64) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (int64_t*)chpl_arr; - parquet::Int64Reader* reader = - static_cast(column_reader.get()); - startIdx -= reader->Skip(startIdx); - - while (reader->HasNext() && arrayIdx < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &chpl_ptr[arrayIdx], &values_read); - // if values_read is 0, that means that it was an empty seg - if (values_read != 0) { - arrayIdx++; - } - i++; - } - } else if(lty == ARROWINT32 || lty == ARROWUINT32) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (int64_t*)chpl_arr; - parquet::Int32Reader* reader = - static_cast(column_reader.get()); - startIdx -= reader->Skip(startIdx); - - int32_t tmp; - while (reader->HasNext() && arrayIdx < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &tmp, &values_read); - // if values_read is 0, that means that it was an empty seg - if (values_read != 0) { - chpl_ptr[arrayIdx] = (int64_t)tmp; - arrayIdx++; - } - i++; - } - } else if (lty == ARROWSTRING) { - 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(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 && definition_level == 3) { - for(int j = 0; j < value.len; j++) { - chpl_ptr[i] = value.ptr[j]; - i++; - } - i++; // skip one space so the strings are null terminated with a 0 - } - } - } else if(lty == ARROWBOOLEAN) { - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (bool*)chpl_arr; - parquet::BoolReader* reader = - static_cast(column_reader.get()); - startIdx -= reader->Skip(startIdx); - - while (reader->HasNext() && arrayIdx < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &chpl_ptr[arrayIdx], &values_read); - // if values_read is 0, that means that it was an empty seg - if (values_read != 0) { - arrayIdx++; - } - i++; - } - } else if(lty == ARROWFLOAT) { - // convert to simpler single batch to sidestep this seemingly architecture dependent (see issue #3234) - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (double*)chpl_arr; - parquet::FloatReader* reader = - static_cast(column_reader.get()); - - float tmp; - while (reader->HasNext() && arrayIdx < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &tmp, &values_read); - // if values_read is 0, that means that it was a null value or empty seg - if (values_read != 0) { - chpl_ptr[arrayIdx] = (double) tmp; - arrayIdx++; - } - else { - // check if nan otherwise it's an empty seg - if (definition_level == 2) { - chpl_ptr[arrayIdx] = NAN; - arrayIdx++; - } - } - i++; - } - } else if(lty == ARROWDOUBLE) { - // convert to simpler single batch to sidestep this seemingly architecture dependent (see issue #3234) - int16_t definition_level; // nullable type and only reading single records in batch - auto chpl_ptr = (double*)chpl_arr; - parquet::DoubleReader* reader = - static_cast(column_reader.get()); - - while (reader->HasNext() && arrayIdx < numElems) { - (void)reader->ReadBatch(1, &definition_level, nullptr, &chpl_ptr[arrayIdx], &values_read); - // if values_read is 0, that means that it was a null value or empty seg - if (values_read != 0) { - arrayIdx++; - } - else { - // check if nan otherwise it's an empty seg - if (definition_level == 2) { - chpl_ptr[arrayIdx] = NAN; - arrayIdx++; - } - } - i++; - } - } - } - return 0; - } - return ARROWERROR; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -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); - auto offsets = (int64_t*)chpl_offsets; - int64_t byteSize = 0; - - if(ty == ARROWSTRING) { - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - int64_t i = 0; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - int64_t idx; - idx = file_metadata -> schema() -> ColumnIndex(colname); - - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - column_reader = row_group_reader->Column(idx); - - int16_t definition_level; - parquet::ByteArrayReader* ba_reader = - static_cast(column_reader.get()); - - int64_t numRead = 0; - - int totalProcessed = 0; - std::vector values(batchSize); - while (ba_reader->HasNext() && totalProcessed < numElems) { - if((numElems - totalProcessed) < batchSize) // adjust batchSize if needed - batchSize = numElems - totalProcessed; - std::vector definition_levels(batchSize,-1); - (void)ba_reader->ReadBatch(batchSize, definition_levels.data(), nullptr, values.data(), &values_read); - totalProcessed += values_read; - int j = 0; - int numProcessed = 0; - while(j < batchSize) { - if(definition_levels[j] == 1 || definition_levels[j] == 3) { - offsets[i] = values[numProcessed].len + 1; - byteSize += values[numProcessed].len + 1; - numProcessed++; - i++; - } else if(definition_levels[j] == 0) { - offsets[i] = 1; - byteSize+=1; - i++; - } else { - j = batchSize; // exit condition - } - j++; - } - } - } - return byteSize; - } - return ARROWERROR; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int64_t cpp_getStringListColumnNumBytes(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 - if (ty == ARROWLIST) { // get the type of the list so we can verify it is ARROWSTRING - dty = cpp_getListType(filename, colname, errMsg); - } - else { - dty = ty; - } - auto offsets = (int64_t*)chpl_offsets; - int64_t byteSize = 0; - - if(dty == ARROWSTRING) { - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - int64_t i = 0; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - int64_t idx; - if (ty == ARROWLIST) { - idx = file_metadata -> schema() -> group_node() -> FieldIndex(colname); - } else { - idx = file_metadata -> schema() -> ColumnIndex(colname); - } - - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - column_reader = row_group_reader->Column(idx); - - int16_t definition_level; - parquet::ByteArrayReader* ba_reader = - static_cast(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; - } - i++; - } - } - } - return byteSize; - } - return ARROWERROR; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int64_t cpp_getListColumnSize(const char* filename, const char* colname, void* chpl_seg_sizes, int64_t numElems, int64_t startIdx, char** errMsg) { - try { - int64_t ty = cpp_getType(filename, colname, errMsg); - auto seg_sizes = (int64_t*)chpl_seg_sizes; - int64_t listSize = 0; - - if (ty == ARROWLIST){ - int64_t lty = cpp_getListType(filename, colname, errMsg); - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - auto idx = file_metadata -> schema() -> group_node() -> FieldIndex(colname); - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - - int64_t i = 0; - int64_t vct = 0; - int64_t seg_size = 0; - int64_t off = 0; - bool first = true; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - column_reader = row_group_reader->Column(idx); - int16_t definition_level; - int16_t rep_lvl; - - if(lty == ARROWINT64 || lty == ARROWUINT64) { - parquet::Int64Reader* int_reader = - static_cast(column_reader.get()); - - while (int_reader->HasNext()) { - int64_t value; - (void)int_reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if (values_read == 0 || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if (values_read != 0 && !int_reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } else if(lty == ARROWINT32 || lty == ARROWUINT32) { - parquet::Int32Reader* int_reader = - static_cast(column_reader.get()); - - while (int_reader->HasNext()) { - int32_t value; - (void)int_reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if (values_read == 0 || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if (values_read != 0 && !int_reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } else if (lty == ARROWSTRING) { - parquet::ByteArrayReader* reader = - static_cast(column_reader.get()); - - while (reader->HasNext()) { - parquet::ByteArray value; - (void)reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if (values_read == 0 || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if (values_read != 0 && !reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } else if(lty == ARROWBOOLEAN) { - parquet::BoolReader* bool_reader = - static_cast(column_reader.get()); - - while (bool_reader->HasNext()) { - bool value; - (void)bool_reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if (values_read == 0 || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if (values_read != 0 && !bool_reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } else if (lty == ARROWFLOAT) { - parquet::FloatReader* float_reader = - static_cast(column_reader.get()); - - int64_t numRead = 0; - while (float_reader->HasNext()) { - float value; - (void)float_reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if ((values_read == 0 && definition_level != 2) || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0 || (values_read == 0 && definition_level == 2)) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if ((values_read != 0 || (values_read == 0 && definition_level == 2)) && !float_reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } else if(lty == ARROWDOUBLE) { - parquet::DoubleReader* dbl_reader = - static_cast(column_reader.get()); - - while (dbl_reader->HasNext()) { - double value; - (void)dbl_reader->ReadBatch(1, &definition_level, &rep_lvl, &value, &values_read); - if ((values_read == 0 && definition_level != 2) || (!first && rep_lvl == 0)) { - seg_sizes[i] = seg_size; - i++; - seg_size = 0; - } - if (values_read != 0 || (values_read == 0 && definition_level == 2)) { - seg_size++; - vct++; - if (first) { - first = false; - } - } - if ((values_read != 0 || (values_read == 0 && definition_level == 2)) && !dbl_reader->HasNext()){ - seg_sizes[i] = seg_size; - } - } - } - } - return vct; - } - return ARROWERROR; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -extern "C" { - int c_readStrColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t batchSize, char** errMsg) { - return cpp_readStrColumnByName(filename, chpl_arr, colname, numElems, batchSize, errMsg); - } - - int c_readAllCols(const char* filename, void** chpl_arrs, int* types, - bool* where_null_chpl, int64_t numElems, - int64_t startIdx, int64_t batchSize, - chplEnum_t nullMode, char** errMsg) { - return akcpp::readAllCols(filename, chpl_arrs, types, where_null_chpl, - numElems, startIdx, batchSize, nullMode, - errMsg); - } - - int c_readColumnByName(const char* filename, void* chpl_arr, - bool* where_null_chpl, const char* colname, - int64_t numElems, int64_t startIdx, int64_t batchSize, - int64_t byteLength, bool hasNonFloatNulls, - char** errMsg) { - return cpp_readColumnByName(filename, chpl_arr, where_null_chpl, colname, - numElems, startIdx, batchSize, byteLength, - hasNonFloatNulls, errMsg); - } - - int c_readListColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg) { - return cpp_readListColumnByName(filename, chpl_arr, colname, numElems, startIdx, batchSize, 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) { - return cpp_getListColumnSize(filename, colname, chpl_seg_sizes, numElems, startIdx, errMsg); - } - - int64_t c_getStringListColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg) { - return cpp_getStringListColumnNumBytes(filename, colname, chpl_offsets, numElems, startIdx, batchSize, errMsg); - } -} diff --git a/src/parquet/ReadParquet.h b/src/parquet/ReadParquet.h deleted file mode 100644 index 23737ef55be..00000000000 --- a/src/parquet/ReadParquet.h +++ /dev/null @@ -1,95 +0,0 @@ -#ifndef READ_PARQUET_H -#define READ_PARQUET_H - -#include -#include - -#include "UtilParquet.h" - - -// Wrap functions in C extern if compiling C++ object file -#ifdef __cplusplus -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace akcpp { - // Engin: We could consider exposing this struct to Chapel. Note that in - // ParquetMsg.chpl there is a somewhat similar type. However, this one is for a - // single column read - // Based on what we decide, the implementation of this struct can be mode more - // C++-like or could stay more C-like. - struct ColReadOp { - void* chpl_arr; - int64_t *startIdx; - std::shared_ptr column_reader; - bool hasNonFloatNulls; - chplEnum_t nullMode; - int64_t row_idx; - int64_t numElems; - int64_t batchSize; - bool* where_null_chpl; - const parquet::ColumnDescriptor* col_info; - - template - int64_t read(); - - template - int64_t _readShortIntegral(); - }; - - int readAllCols(const char* filename, void** chpl_arrs, int* types, - bool* where_null_chpl, int64_t numElems, int64_t startIdx, - int64_t batchSize, chplEnum_t nullMode, char** errMsg); -} - - -extern "C" { -#endif - int c_readStrColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t batchSize, char** errMsg); - - int cpp_readStrColumnByName(const char* filename, void* chpl_arr, const char* colname, int64_t numElems, int64_t batchSize, char** errMsg); - - int c_readAllCols(const char* filename, void** chpl_arrs, int* types, - bool* where_null_chpl, int64_t numElems, - int64_t startIdx, int64_t batchSize, - chplEnum_t nullMode, char** errMsg); - - int c_readColumnByName(const char* filename, void* chpl_arr, bool* where_null_chpl, - const char* colname, int64_t numElems, int64_t startIdx, - int64_t batchSize, int64_t byteLength, bool hasNonFloatNulls, char** errMsg); - int cpp_readColumnByName(const char* filename, void* chpl_arr, bool* where_null_chpl, - const char* colname, int64_t numElems, int64_t startIdx, - int64_t batchSize, int64_t byteLength, bool hasNonFloatNulls, char** errMsg); - - int c_readListColumnByName(const char* filename, void* chpl_arr, - const char* colname, int64_t numElems, - int64_t startIdx, int64_t batchSize, char** errMsg); - int cpp_readListColumnByName(const char* filename, void* chpl_arr, - 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, 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); - int64_t cpp_getListColumnSize(const char* filename, const char* colname, - void* chpl_seg_sizes, int64_t numElems, int64_t startIdx, char** errMsg); - - int64_t c_getStringListColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg); - int64_t cpp_getStringListColumnNumBytes(const char* filename, const char* colname, void* chpl_offsets, int64_t numElems, int64_t startIdx, int64_t batchSize, char** errMsg); - -#ifdef __cplusplus -} -#endif -#endif //READ_PARQUET_H diff --git a/src/parquet/SharedEnums.h b/src/parquet/SharedEnums.h deleted file mode 100644 index 253ba8bd652..00000000000 --- a/src/parquet/SharedEnums.h +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef SHARED_ENUMS_H -#define SHARED_ENUMS_H - -#include "../ParquetSharedEnums.chpl" - -#endif // SHARED_ENUMS_H diff --git a/src/parquet/UtilParquet.cpp b/src/parquet/UtilParquet.cpp deleted file mode 100644 index f72d4953b8d..00000000000 --- a/src/parquet/UtilParquet.cpp +++ /dev/null @@ -1,903 +0,0 @@ -#include "UtilParquet.h" - - /* - Arrow Error Helpers - ------------------- - Arrow provides PARQUETASSIGNORTHROW and other similar macros - to help with error handling, but since we are doing something - unique (passing back the error message to Chapel to be displayed), - these helpers are similar to the provided macros but matching our - functionality. - */ - - // The `ARROWRESULT_OK` macro should be used when trying to - // assign the result of an Arrow/Parquet function to a value that can - // potentially throw an error, so the argument `cmd` is the Arrow - // command to execute and `res` is the desired variable to store the - // result -#define ARROWRESULT_OK(cmd, res) \ - { \ - auto result = cmd; \ - if(!result.ok()) { \ - *errMsg = strdup(result.status().message().c_str()); \ - return ARROWERROR; \ - } \ - res = std::move(result.ValueOrDie()); \ - } - - // The `ARROWSTATUS_OK` macro should be used when calling an - // Arrow/Parquet function that returns a status. The `cmd` - // argument should be the Arrow function to execute. -#define ARROWSTATUS_OK(cmd) \ - if(!check_status_ok(cmd, errMsg)) \ - return ARROWERROR; - -static std::map> globalFiles; -static std::map> globalRowGroupReaders; -static std::map> globalColumnReaders; - -bool check_status_ok(arrow::Status status, char** errMsg) { - if(!status.ok()) { - *errMsg = strdup(status.message().c_str()); - return false; - } - return true; -} - -/* - C++ functions - ------------- - These C++ functions are used to call into the Arrow library - and are then called to by their corresponding C functions to - allow interoperability with Chapel. This means that all of the - C++ functions must return types that are C compatible. -*/ - -int64_t cpp_getNumRows(const char* filename, char** errMsg) { - try { - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - - return reader -> parquet_reader() -> metadata() -> num_rows(); - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -static int getSchema(const char* filename, - std::shared_ptr* out, - char** errMsg) { - - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, - arrow::default_memory_pool()), - infile); - - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - ARROWSTATUS_OK(reader->GetSchema(out)); - - return 0; -} - -int64_t cpp_getNumCols(const char* filename, char** errMsg) { - std::shared_ptr sc; - - if (getSchema(filename, &sc, errMsg) == ARROWERROR) { - return ARROWERROR; - } - - return sc->num_fields(); -} - -int cpp_getPrecision(const char* filename, const char* colname, char** errMsg) { - try { - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - - std::shared_ptr sc; - std::shared_ptr* out = ≻ - ARROWSTATUS_OK(reader->GetSchema(out)); - - int idx = sc -> GetFieldIndex(colname); - - const auto& decimal_type = static_cast(*sc->field(idx)->type()); - const int64_t precision = decimal_type.precision(); - - return precision; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -static int getTypeOfFieldAtIdx(std::shared_ptr sc, - int idx) { - auto myType = sc -> field(idx) -> type(); - - if(myType->id() == arrow::Type::INT64) - return ARROWINT64; - else if(myType->id() == arrow::Type::INT32 || myType->id() == arrow::Type::INT16) - return ARROWINT32; // int16 is logical type, stored as int32 - else if(myType->id() == arrow::Type::UINT64) - return ARROWUINT64; - else if(myType->id() == arrow::Type::UINT32 || - myType->id() == arrow::Type::UINT16) - return ARROWUINT32; // uint16 is logical type, stored as uint32 - else if(myType->id() == arrow::Type::TIMESTAMP) - return ARROWTIMESTAMP; - else if(myType->id() == arrow::Type::BOOL) - return ARROWBOOLEAN; - else if(myType->id() == arrow::Type::STRING || - myType->id() == arrow::Type::BINARY || - myType->id() == arrow::Type::LARGE_STRING) - return ARROWSTRING; - else if(myType->id() == arrow::Type::FLOAT) - return ARROWFLOAT; - else if(myType->id() == arrow::Type::DOUBLE) - return ARROWDOUBLE; - else if(myType->id() == arrow::Type::LIST) - return ARROWLIST; - else if(myType->id() == arrow::Type::DECIMAL) - return ARROWDECIMAL; - else { - return ARROWERROR; - } -} - -int cpp_getAllTypes(const char* filename, int* types_out, char** errMsg) { - assert(types_out); // this has to be allocated by the caller - - std::shared_ptr sc; - - try { - getSchema(filename, &sc, errMsg); - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } - - const int num_cols = sc->num_fields(); - for(int i=0 ; i infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - - std::shared_ptr sc; - std::shared_ptr* out = ≻ - ARROWSTATUS_OK(reader->GetSchema(out)); - - int idx = sc -> GetFieldIndex(colname); - // Since this doesn't actually throw a Parquet error, we have to generate - // our own error message for this case - if(idx == -1) { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Dataset: " + dname + " does not exist in file: " + filename; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - auto myType = sc -> field(idx) -> type(); - - if(myType->id() == arrow::Type::INT64) - return ARROWINT64; - else if(myType->id() == arrow::Type::INT32 || myType->id() == arrow::Type::INT16) - return ARROWINT32; // int16 is logical type, stored as int32 - else if(myType->id() == arrow::Type::UINT64) - return ARROWUINT64; - else if(myType->id() == arrow::Type::UINT32 || - myType->id() == arrow::Type::UINT16) - return ARROWUINT32; // uint16 is logical type, stored as uint32 - else if(myType->id() == arrow::Type::TIMESTAMP) - return ARROWTIMESTAMP; - else if(myType->id() == arrow::Type::BOOL) - return ARROWBOOLEAN; - else if(myType->id() == arrow::Type::STRING || - myType->id() == arrow::Type::BINARY || - myType->id() == arrow::Type::LARGE_STRING) - return ARROWSTRING; - else if(myType->id() == arrow::Type::FLOAT) - return ARROWFLOAT; - else if(myType->id() == arrow::Type::DOUBLE) - return ARROWDOUBLE; - else if(myType->id() == arrow::Type::LIST) - return ARROWLIST; - else if(myType->id() == arrow::Type::DECIMAL) - return ARROWDECIMAL; - else { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Unsupported type on column: " + dname + " in " + filename; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_getListType(const char* filename, const char* colname, char** errMsg) { - try { - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - - std::shared_ptr sc; - std::shared_ptr* out = ≻ - ARROWSTATUS_OK(reader->GetSchema(out)); - - int idx = sc -> GetFieldIndex(colname); - // Since this doesn't actually throw a Parquet error, we have to generate - // our own error message for this case - if(idx == -1) { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - auto myType = sc -> field(idx) -> type(); - - if (myType->id() == arrow::Type::LIST) { - if (myType->num_fields() != 1) { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Column " + dname + " in " + fname + " cannot be read by Arkouda."; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - else { - // fields returns a vector of fields, but here we are expecting lists so should only contain 1 item here - auto field = myType->fields()[0]; - auto f_type = field->type(); - if(f_type->id() == arrow::Type::INT64) - return ARROWINT64; - else if(f_type->id() == arrow::Type::INT32 || f_type->id() == arrow::Type::INT16) - return ARROWINT32; - else if(f_type->id() == arrow::Type::UINT64) - return ARROWUINT64; - else if(f_type->id() == arrow::Type::UINT32 || f_type->id() == arrow::Type::UINT16) - return ARROWUINT32; - else if(f_type->id() == arrow::Type::TIMESTAMP) - return ARROWTIMESTAMP; - else if(f_type->id() == arrow::Type::BOOL) - return ARROWBOOLEAN; - else if(f_type->id() == arrow::Type::STRING || - f_type->id() == arrow::Type::BINARY || - f_type->id() == arrow::Type::LARGE_STRING) // Verify that this is functional as expected - return ARROWSTRING; - else if(f_type->id() == arrow::Type::FLOAT) - return ARROWFLOAT; - else if(f_type->id() == arrow::Type::DOUBLE) - return ARROWDOUBLE; - else { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Unsupported type on column: " + dname + " in " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - } - } - else { - std::string fname(filename); - std::string dname(colname); - std::string msg = "Column " + dname + " in " + fname + " is not a List"; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int64_t cpp_getStringColumnNullIndices(const char* filename, const char* colname, void* chpl_nulls, char** errMsg) { - try { - int64_t ty = cpp_getType(filename, colname, errMsg); - auto null_indices = (int64_t*)chpl_nulls; - int64_t byteSize = 0; - - if(ty == ARROWSTRING) { - std::unique_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - - std::shared_ptr file_metadata = parquet_reader->metadata(); - int num_row_groups = file_metadata->num_row_groups(); - - int64_t i = 0; - for (int r = 0; r < num_row_groups; r++) { - std::shared_ptr row_group_reader = - parquet_reader->RowGroup(r); - - int64_t values_read = 0; - - std::shared_ptr column_reader; - - auto idx = file_metadata -> schema() -> ColumnIndex(colname); - - if(idx < 0) { - std::string dname(colname); - std::string fname(filename); - std::string msg = "Dataset: " + dname + " does not exist in file: " + fname; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - column_reader = row_group_reader->Column(idx); - int16_t definition_level; - parquet::ByteArrayReader* ba_reader = - static_cast(column_reader.get()); - - while (ba_reader->HasNext()) { - parquet::ByteArray value; - (void)ba_reader->ReadBatch(1, &definition_level, nullptr, &value, &values_read); - if(values_read == 0) - null_indices[i] = 1; - i++; - } - } - return 0; - } - return ARROWERROR; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -// configure the schema for a multicolumn file -std::shared_ptr SetupSchema(void* column_names, void * objTypes, void* datatypes, int64_t colnum) { - parquet::schema::NodeVector fields; - auto cname_ptr = (char**)column_names; - auto dtypes_ptr = (int64_t*) datatypes; - auto objType_ptr = (int64_t*) objTypes; - for (int64_t i = 0; i < colnum; i++){ - if(dtypes_ptr[i] == ARROWINT64) { - if (objType_ptr[i] == SEGARRAY){ - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } else { - fields.push_back(parquet::schema::PrimitiveNode::Make(cname_ptr[i], parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::NONE)); - } - } else if(dtypes_ptr[i] == ARROWUINT64) { - if (objType_ptr[i] == SEGARRAY){ - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::UINT_64); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } else { - fields.push_back(parquet::schema::PrimitiveNode::Make(cname_ptr[i], parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::UINT_64)); - } - } else if(dtypes_ptr[i] == ARROWBOOLEAN) { - if (objType_ptr[i] == SEGARRAY){ - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } else { - fields.push_back(parquet::schema::PrimitiveNode::Make(cname_ptr[i], parquet::Repetition::REQUIRED, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE)); - } - } else if(dtypes_ptr[i] == ARROWDOUBLE) { - if (objType_ptr[i] == SEGARRAY) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::DOUBLE, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } else { - fields.push_back(parquet::schema::PrimitiveNode::Make(cname_ptr[i], parquet::Repetition::REQUIRED, parquet::Type::DOUBLE, parquet::ConvertedType::NONE)); - } - } else if(dtypes_ptr[i] == ARROWSTRING) { - if (objType_ptr[i] == SEGARRAY) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::BYTE_ARRAY, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } else { - fields.push_back(parquet::schema::PrimitiveNode::Make(cname_ptr[i], parquet::Repetition::OPTIONAL, parquet::Type::BYTE_ARRAY, parquet::ConvertedType::NONE)); - } - } - } - return std::static_pointer_cast( - parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); -} - -int cpp_createEmptyListParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg) { - try { - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - PARQUET_ASSIGN_OR_THROW(out_file, FileClass::Open(filename)); - - parquet::schema::NodeVector fields; - if (dtype == ARROWINT64) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWUINT64) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::UINT_64); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWBOOLEAN) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWDOUBLE) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::DOUBLE, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_createEmptyParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg) { - try { - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - PARQUET_ASSIGN_OR_THROW(out_file, FileClass::Open(filename)); - - parquet::schema::NodeVector fields; - if(dtype == ARROWINT64) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::NONE)); - else if(dtype == ARROWUINT64) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::UINT_64)); - else if(dtype == ARROWBOOLEAN) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE)); - else if(dtype == ARROWDOUBLE) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::DOUBLE, parquet::ConvertedType::NONE)); - else if(dtype == ARROWSTRING) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::OPTIONAL, parquet::Type::BYTE_ARRAY, parquet::ConvertedType::NONE)); - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_appendColumnToParquet(const char* filename, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t dtype, int64_t compression, - char** errMsg) { - try { - if (chpl_arr == NULL){ - // early out to prevent bad memory access - return 0; - } - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - // Use threads for case when reading a table with many columns - reader->set_use_threads(true); - -#if ARROW_VERSION_MAJOR < 24 - std::shared_ptr table; - std::shared_ptr* hold_table = &table; - ARROWSTATUS_OK(reader->ReadTable(hold_table)); -#else - std::shared_ptr table; - ARROWRESULT_OK(reader->ReadTable(), table); -#endif - - - arrow::ArrayVector arrays; - std::shared_ptr values; - auto chunk_type = arrow::int64(); - if(dtype == ARROWINT64) { - chunk_type = arrow::int64(); - arrow::Int64Builder builder; - auto chpl_ptr = (int64_t*)chpl_arr; - ARROWSTATUS_OK(builder.AppendValues(chpl_ptr, numelems, nullptr)) - ARROWSTATUS_OK(builder.Finish(&values)); - } else if(dtype == ARROWUINT64) { - chunk_type = arrow::uint64(); - arrow::UInt64Builder builder; - auto chpl_ptr = (uint64_t*)chpl_arr; - ARROWSTATUS_OK(builder.AppendValues(chpl_ptr, numelems, nullptr)) - ARROWSTATUS_OK(builder.Finish(&values)); - } else if(dtype == ARROWBOOLEAN) { - chunk_type = arrow::boolean(); - arrow::BooleanBuilder builder; - auto chpl_ptr = (uint8_t*)chpl_arr; - ARROWSTATUS_OK(builder.AppendValues(chpl_ptr, numelems, nullptr)) - ARROWSTATUS_OK(builder.Finish(&values)); - } else if(dtype == ARROWSTRING) { - chunk_type = arrow::utf8(); - arrow::StringBuilder builder; - auto chpl_ptr = (uint8_t*)chpl_arr; - int64_t j = 0; - for(int64_t i = 0; i < numelems; i++) { - std::string tmp_str = ""; - while(chpl_ptr[j] != 0x00) { - tmp_str += chpl_ptr[j++]; - } - j++; - - auto const status = builder.Append(tmp_str); - if (status.IsCapacityError()) { - // Reached current chunk's capacity limit, so start a new one... - ARROWSTATUS_OK(builder.Finish(&values)); - arrays.push_back(values); - values.reset(); - builder.Reset(); - - // ...with this string as its first item. - ARROWSTATUS_OK(builder.Append(tmp_str)); - } else { - ARROWSTATUS_OK(status); - } - } - ARROWSTATUS_OK(builder.Finish(&values)); - } else if(dtype == ARROWDOUBLE) { - chunk_type = arrow::float64(); - arrow::DoubleBuilder builder; - auto chpl_ptr = (double*)chpl_arr; - ARROWSTATUS_OK(builder.AppendValues(chpl_ptr, numelems, nullptr)) - ARROWSTATUS_OK(builder.Finish(&values)); - } else { - std::string msg = "Unrecognized Parquet dtype"; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - arrays.push_back(values); - - std::shared_ptr chunk_sh_ptr; - ARROWRESULT_OK(arrow::ChunkedArray::Make({arrays}, chunk_type), chunk_sh_ptr); - - auto newField = arrow::field(dsetname, chunk_type); - std::shared_ptr fin_table; - ARROWRESULT_OK(table -> AddColumn(0, newField, chunk_sh_ptr), fin_table); - - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - ARROWRESULT_OK(FileClass::Open(filename), out_file); - ARROWSTATUS_OK(parquet::arrow::WriteTable(*fin_table, arrow::default_memory_pool(), out_file, numelems)); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -const char* cpp_getVersionInfo(void) { - return strdup(arrow::GetBuildInfo().version_string.c_str()); -} - -int cpp_getDatasetNames(const char* filename, char** dsetResult, bool readNested, char** errMsg) { - try { - std::shared_ptr infile; - ARROWRESULT_OK(arrow::io::ReadableFile::Open(filename, arrow::default_memory_pool()), - infile); - std::unique_ptr reader; -#if ARROW_VERSION_MAJOR < 19 - ARROWSTATUS_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader)); -#else - ARROWRESULT_OK(parquet::arrow::OpenFile(infile, arrow::default_memory_pool()), reader); -#endif - - std::shared_ptr sc; - std::shared_ptr* out = ≻ - ARROWSTATUS_OK(reader->GetSchema(out)); - - std::string fields = ""; - bool first = true; - - for(int i = 0; i < sc->num_fields(); i++) { - // only add fields of supported types - if(sc->field(i)->type()->id() == arrow::Type::INT64 || - sc->field(i)->type()->id() == arrow::Type::INT32 || - sc->field(i)->type()->id() == arrow::Type::INT16 || - sc->field(i)->type()->id() == arrow::Type::UINT64 || - sc->field(i)->type()->id() == arrow::Type::UINT32 || - sc->field(i)->type()->id() == arrow::Type::UINT16 || - sc->field(i)->type()->id() == arrow::Type::TIMESTAMP || - sc->field(i)->type()->id() == arrow::Type::BOOL || - sc->field(i)->type()->id() == arrow::Type::STRING || - sc->field(i)->type()->id() == arrow::Type::BINARY || - sc->field(i)->type()->id() == arrow::Type::FLOAT || - sc->field(i)->type()->id() == arrow::Type::DOUBLE || - (sc->field(i)->type()->id() == arrow::Type::LIST && readNested) || - sc->field(i)->type()->id() == arrow::Type::DECIMAL || - sc->field(i)->type()->id() == arrow::Type::LARGE_STRING - ) { - if(!first) - fields += ("," + sc->field(i)->name()); - else - fields += (sc->field(i)->name()); - first = false; - } else if (sc->field(i)->type()->id() == arrow::Type::LIST && !readNested) { - continue; - } else { - std::string fname(filename); - std::string dname(sc->field(i)->ToString()); - std::string msg = "Unsupported type on column: " + dname + " in " + filename; - *errMsg = strdup(msg.c_str()); - return ARROWERROR; - } - } - *dsetResult = strdup(fields.c_str()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -void cpp_free_string(void* ptr) { - free(ptr); -} - -void cpp_openFile(const char* filename, int64_t idx) { - std::shared_ptr parquet_reader = - parquet::ParquetFileReader::OpenFile(filename, false); - globalFiles[idx] = parquet_reader; -} - -void cpp_createRowGroupReader(int64_t rowGroup, int64_t readerIdx) { - std::shared_ptr row_group_reader = - globalFiles[readerIdx]->RowGroup(rowGroup); - globalRowGroupReaders[readerIdx] = row_group_reader; -} - -void cpp_createColumnReader(const char* colname, int64_t readerIdx) { - - std::shared_ptr file_metadata = globalFiles[readerIdx]->metadata(); - auto idx = file_metadata -> schema() -> ColumnIndex(colname); - - std::shared_ptr column_reader; - column_reader = globalRowGroupReaders[readerIdx]->Column(idx); - globalColumnReaders[readerIdx] = column_reader; -} - - - -int cpp_getNumRowGroups(int64_t readerIdx) { - std::shared_ptr file_metadata = globalFiles[readerIdx]->metadata(); - return file_metadata->num_row_groups(); -} - -void cpp_freeMapValues(void* row) { - parquet::ByteArray* string_values = - static_cast(row); - free(string_values); - globalColumnReaders.clear(); - globalRowGroupReaders.clear(); - globalFiles.clear(); -} - -int cpp_readParquetColumnChunks(const char* filename, int64_t batchSize, int64_t numElems, - int64_t readerIdx, int64_t* numRead, - void** outData, bool* containsNulls, char** errMsg) { - try { - auto reader = static_cast(globalColumnReaders[readerIdx].get()); - parquet::ByteArray* string_values = - (parquet::ByteArray*)malloc(numElems*sizeof(parquet::ByteArray)); - std::vector definition_level(batchSize); - int64_t values_read = 0; - int64_t total_read = 0; - while(reader->HasNext() && total_read < numElems) { - if((numElems - total_read) < batchSize) - batchSize = numElems - total_read; - // adding 1 to definition level, since the first value indicates if null values - (void)reader->ReadBatch(batchSize, definition_level.data(), nullptr, string_values + total_read, &values_read); - for(int i = 0; i < values_read; i++) { - if(definition_level[i] == 0) - *containsNulls = true; - } - total_read += values_read; - } - *numRead = total_read; - *outData = (void*)string_values; - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -/* - C functions - ----------- - These C functions provide no functionality, since the C++ - Arrow library is being used, they merely call the C++ functions - to allow Chapel to call the C++ functions through C interoperability. - Each Arrow function must have a corresponding C function if wished - to be called by Chapel. -*/ - -extern "C" { - int64_t c_getNumRows(const char* filename, char** errMsg) { - return cpp_getNumRows(filename, errMsg); - } - - int64_t c_getNumCols(const char* filename, char** errMsg) { - return cpp_getNumCols(filename, errMsg); - } - - int c_getType(const char* filename, const char* colname, char** errMsg) { - return cpp_getType(filename, colname, errMsg); - } - - int c_getAllTypes(const char* filename, int* types_out, char** errMsg) { - return cpp_getAllTypes(filename, types_out, errMsg); - } - - int c_getListType(const char* filename, const char* colname, char** errMsg) { - return cpp_getListType(filename, colname, errMsg); - } - - int c_createEmptyParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg) { - return cpp_createEmptyParquetFile(filename, dsetname, dtype, compression, errMsg); - } - - int c_createEmptyListParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg) { - return cpp_createEmptyListParquetFile(filename, dsetname, dtype, compression, errMsg); - } - - int c_appendColumnToParquet(const char* filename, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t dtype, int64_t compression, - char** errMsg) { - return cpp_appendColumnToParquet(filename, chpl_arr, - dsetname, numelems, - dtype, compression, - errMsg); - } - - int64_t c_getStringColumnNullIndices(const char* filename, const char* colname, void* chpl_nulls, char** errMsg) { - return cpp_getStringColumnNullIndices(filename, colname, chpl_nulls, errMsg); - } - - const char* c_getVersionInfo(void) { - return cpp_getVersionInfo(); - } - - int c_getDatasetNames(const char* filename, char** dsetResult, bool readNested, char** errMsg) { - return cpp_getDatasetNames(filename, dsetResult, readNested, errMsg); - } - - void c_free_string(void* ptr) { - cpp_free_string(ptr); - } - - int c_getPrecision(const char* filename, const char* colname, char** errMsg) { - return cpp_getPrecision(filename, colname, errMsg); - } - - void c_openFile(const char* filename, int64_t idx) { - cpp_openFile(filename, idx); - } - - void c_createRowGroupReader(int64_t rowGroup, int64_t readerIdx) { - return cpp_createRowGroupReader(rowGroup, readerIdx); - } - - void c_createColumnReader(const char* colname, int64_t readerIdx) { - cpp_createColumnReader(colname, readerIdx); - } - - int c_getNumRowGroups(int64_t readerIdx) { - return cpp_getNumRowGroups(readerIdx); - } - - void c_freeMapValues(void* row) { - cpp_freeMapValues(row); - } - - int c_readParquetColumnChunks(const char* filename, int64_t batchSize, int64_t numElems, - int64_t readerIdx, int64_t* numRead, - void** outData, bool* containsNulls, char** errMsg) { - return cpp_readParquetColumnChunks(filename, batchSize, numElems, readerIdx, - numRead, outData, containsNulls, errMsg); - } -} diff --git a/src/parquet/UtilParquet.h b/src/parquet/UtilParquet.h deleted file mode 100644 index 61388368733..00000000000 --- a/src/parquet/UtilParquet.h +++ /dev/null @@ -1,166 +0,0 @@ -#ifndef UTIL_PARQUET_H -#define UTIL_PARQUET_H - -#include -#include -#include "SharedEnums.h" - -#define chplEnum_t int64_t - -// Wrap functions in C extern if compiling C++ object file -#ifdef __cplusplus -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - - -std::shared_ptr SetupSchema(void* column_names, void * objTypes, void* datatypes, int64_t colnum); - - -extern "C" { -#endif - -#define ARROWINT64 0 -#define ARROWINT32 1 -#define ARROWUINT64 2 -#define ARROWUINT32 3 -#define ARROWBOOLEAN 4 -#define ARROWFLOAT 5 -#define ARROWDOUBLE 7 -#define ARROWTIMESTAMP ARROWINT64 -#define ARROWSTRING 6 -#define ARROWLIST 8 -#define ARROWDECIMAL 9 -#define ARROWERROR -1 - -#define ARRAYVIEW 0 // not currently used, but included for continuity with Chapel -#define PDARRAY 1 -#define STRINGS 2 -#define SEGARRAY 3 - -// compression mappings -#define SNAPPY_COMP 1 -#define GZIP_COMP 2 -#define BROTLI_COMP 3 -#define ZSTD_COMP 4 -#define LZ4_COMP 5 - - - typedef struct { - uint32_t len; - const uint8_t* ptr; - } MyByteArray; - - void c_openFile(const char* filename, int64_t idx); - void cpp_openFile(const char* filename, int64_t idx); - - void c_createRowGroupReader(int64_t rowGroup, int64_t readerIdx); - void cpp_createRowGroupReader(int64_t rowGroup, int64_t readerIdx); - - void c_createColumnReader(const char* colname, int64_t readerIdx); - void cpp_createColumnReader(const char* colname, int64_t readerIdx); - - int c_getNumRowGroups(int64_t readerIdx); - int cpp_getNumRowGroups(int64_t readerIdx); - - // Each C++ function contains the actual implementation of the - // functionality, and there is a corresponding C function that - // Chapel can call into through C interoperability, since there - // is no C++ interoperability supported in Chapel today. - int64_t c_getNumRows(const char*, char** errMsg); - int64_t cpp_getNumRows(const char*, char** errMsg); - - int64_t c_getNumCols(const char* filename, char** errMsg); - int64_t cpp_getNumCols(const char* filename, char** errMsg); - - int c_getAllTypes(const char* filename, int* types_out, char** errMsg); - - int64_t c_getStringColumnNullIndices(const char* filename, const char* colname, void* chpl_nulls, char** errMsg); - int64_t cpp_getStringColumnNullIndices(const char* filename, const char* colname, void* chpl_nulls, char** errMsg); - - int c_getType(const char* filename, const char* colname, char** errMsg); - int cpp_getType(const char* filename, const char* colname, char** errMsg); - - int c_getListType(const char* filename, const char* colname, char** errMsg); - int cpp_getListType(const char* filename, const char* colname, char** errMsg); - - int cpp_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int c_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, char** errMsg); - - int c_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int cpp_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - - int c_createEmptyListParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg); - int cpp_createEmptyListParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg); - - int c_writeListColumnToParquet(const char* filename, void* chpl_offsets, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int cpp_writeListColumnToParquet(const char* filename, void* chpl_offsets, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - - int c_createEmptyParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg); - int cpp_createEmptyParquetFile(const char* filename, const char* dsetname, int64_t dtype, - int64_t compression, char** errMsg); - - int c_appendColumnToParquet(const char* filename, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t dtype, int64_t compression, - char** errMsg); - int cpp_appendColumnToParquet(const char* filename, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t dtype, int64_t compression, - char** errMsg); - - int c_getPrecision(const char* filename, const char* colname, char** errMsg); - int cpp_getPrecision(const char* filename, const char* colname, char** errMsg); - - const char* c_getVersionInfo(void); - const char* cpp_getVersionInfo(void); - - int c_getDatasetNames(const char* filename, char** dsetResult, bool readNested, char** errMsg); - int cpp_getDatasetNames(const char* filename, char** dsetResult, bool readNested, char** errMsg); - - void c_free_string(void* ptr); - void cpp_free_string(void* ptr); - - void c_freeMapValues(void* row); - void cpp_freeMapValues(void* row); - - int c_readParquetColumnChunks(const char* filename, int64_t batchSize, int64_t numElems, - int64_t readerIdx, int64_t* numRead, - void** outData, bool* containsNulls, char** errMsg); - int cpp_readParquetColumnChunks(const char* filename, int64_t batchSize, int64_t numElems, - int64_t readerIdx, int64_t* numRead, - void** outData, bool* containsNulls, char** errMsg); - -#ifdef __cplusplus - bool check_status_ok(arrow::Status status, char** errMsg); -} -#endif -#endif //UTIL_PARQUET_H diff --git a/src/parquet/WriteParquet.cpp b/src/parquet/WriteParquet.cpp deleted file mode 100644 index b8c6a07d393..00000000000 --- a/src/parquet/WriteParquet.cpp +++ /dev/null @@ -1,829 +0,0 @@ -#include "WriteParquet.h" -#include "UtilParquet.h" - - /* - Arrow Error Helpers - ------------------- - Arrow provides PARQUETASSIGNORTHROW and other similar macros - to help with error handling, but since we are doing something - unique (passing back the error message to Chapel to be displayed), - these helpers are similar to the provided macros but matching our - functionality. - */ - - // The `ARROWRESULT_OK` macro should be used when trying to - // assign the result of an Arrow/Parquet function to a value that can - // potentially throw an error, so the argument `cmd` is the Arrow - // command to execute and `res` is the desired variable to store the - // result -#define ARROWRESULT_OK(cmd, res) \ - { \ - auto result = cmd; \ - if(!result.ok()) { \ - *errMsg = strdup(result.status().message().c_str()); \ - return ARROWERROR; \ - } \ - res = result.ValueOrDie(); \ - } - - // The `ARROWSTATUS_OK` macro should be used when calling an - // Arrow/Parquet function that returns a status. The `cmd` - // argument should be the Arrow function to execute. -#define ARROWSTATUS_OK(cmd) \ - if(!check_status_ok(cmd, errMsg)) \ - return ARROWERROR; - -int cpp_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - try { - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - ARROWRESULT_OK(FileClass::Open(filename), out_file); - - parquet::schema::NodeVector fields; - if(dtype == ARROWINT64) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::NONE)); - else if(dtype == ARROWUINT64) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::INT64, parquet::ConvertedType::UINT_64)); - else if(dtype == ARROWBOOLEAN) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE)); - else if(dtype == ARROWDOUBLE) - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::REQUIRED, parquet::Type::DOUBLE, parquet::ConvertedType::NONE)); - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - int64_t i = 0; - int64_t numLeft = numelems; - - if (chpl_arr == NULL) { - // early out to prevent bad memory access - return 0; - } - - if(dtype == ARROWINT64 || dtype == ARROWUINT64) { - auto chpl_ptr = (int64_t*)chpl_arr; - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::Int64Writer* int64_writer = - static_cast(rg_writer->NextColumn()); - - int64_t batchSize = rowGroupSize; - if(numLeft < rowGroupSize) - batchSize = numLeft; - int64_writer->WriteBatch(batchSize, nullptr, nullptr, &chpl_ptr[i]); - numLeft -= batchSize; - i += batchSize; - } - } else if(dtype == ARROWBOOLEAN) { - auto chpl_ptr = (bool*)chpl_arr; - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::BoolWriter* writer = - static_cast(rg_writer->NextColumn()); - - int64_t batchSize = rowGroupSize; - if(numLeft < rowGroupSize) - batchSize = numLeft; - writer->WriteBatch(batchSize, nullptr, nullptr, &chpl_ptr[i]); - numLeft -= batchSize; - i += batchSize; - } - } else if(dtype == ARROWDOUBLE) { - auto chpl_ptr = (double*)chpl_arr; - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::DoubleWriter* writer = - static_cast(rg_writer->NextColumn()); - - int64_t batchSize = rowGroupSize; - if(numLeft < rowGroupSize) - batchSize = numLeft; - writer->WriteBatch(batchSize, nullptr, nullptr, &chpl_ptr[i]); - numLeft -= batchSize; - i += batchSize; - } - } else { - return ARROWERROR; - } - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - try { - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - PARQUET_ASSIGN_OR_THROW(out_file, FileClass::Open(filename)); - - parquet::schema::NodeVector fields; - - fields.push_back(parquet::schema::PrimitiveNode::Make(dsetname, parquet::Repetition::OPTIONAL, parquet::Type::BYTE_ARRAY, parquet::ConvertedType::NONE)); - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - int64_t i = 0; - int64_t numLeft = numelems; - - if(dtype == ARROWSTRING) { - auto chpl_ptr = (uint8_t*)chpl_arr; - auto offsets = (int64_t*)chpl_offsets; - int64_t byteIdx = 0; - int64_t offIdx = 0; - - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::ByteArrayWriter* ba_writer = - static_cast(rg_writer->NextColumn()); - int64_t count = 0; - while(numLeft > 0 && count < rowGroupSize) { - parquet::ByteArray value; - int16_t definition_level = 1; - value.ptr = reinterpret_cast(&chpl_ptr[byteIdx]); - // subtract 1 since we have the null terminator - value.len = offsets[offIdx+1] - offsets[offIdx] - 1; - if (value.len == 0) - definition_level = 0; - ba_writer->WriteBatch(1, &definition_level, nullptr, &value); - numLeft--;count++; - offIdx++; - byteIdx+=offsets[offIdx] - offsets[offIdx-1]; - } - } - } else { - return ARROWERROR; - } - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_writeStrListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_offsets, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - try { - if(dtype == ARROWSTRING) { // check the type here so if it is wrong we don't create a bad file - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - PARQUET_ASSIGN_OR_THROW(out_file, FileClass::Open(filename)); - - parquet::schema::NodeVector fields; - - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::BYTE_ARRAY, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - int64_t i = 0; - int64_t numLeft = numelems; - auto segments = (int64_t*)chpl_segs; - int64_t segIdx = 0; // index into segarray segments - int64_t offIdx = 0; // index into the segstring segments - int64_t valIdx = 0; // index into chpl_arr - - while(numLeft > 0) { // write all local values to the file - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::ByteArrayWriter* ba_writer = - static_cast(rg_writer->NextColumn()); - int64_t count = 0; - while (numLeft > 0 && count < rowGroupSize) { // ensures rowGroupSize maintained - int64_t segmentLength = segments[segIdx+1] - segments[segIdx]; - if (segmentLength > 0) { - auto offsets = (int64_t*)chpl_offsets; - auto chpl_ptr = (uint8_t*)chpl_arr; - for (int64_t x = 0; x < segmentLength; x++){ - int16_t rep_lvl = (x == 0) ? 0 : 1; - int16_t def_lvl = 3; - parquet::ByteArray value; - value.ptr = reinterpret_cast(&chpl_ptr[valIdx]); - value.len = offsets[offIdx+1] - offsets[offIdx] - 1; - ba_writer->WriteBatch(1, &def_lvl, &rep_lvl, &value); - offIdx++; - valIdx+=offsets[offIdx] - offsets[offIdx-1]; - } - } else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - segmentLength = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - ba_writer->WriteBatch(segmentLength, &def_lvl, &rep_lvl, nullptr); - } - segIdx++; - numLeft--;count++; - } - } - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - return 0; - } else { - return ARROWERROR; - } - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_writeListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - try { - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - PARQUET_ASSIGN_OR_THROW(out_file, FileClass::Open(filename)); - - parquet::schema::NodeVector fields; - - // create the list schema. List containing the dtype - if (dtype == ARROWINT64) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWUINT64) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::INT64, parquet::ConvertedType::UINT_64); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWBOOLEAN) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::BOOLEAN, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - else if (dtype == ARROWDOUBLE) { - auto element = parquet::schema::PrimitiveNode::Make("item", parquet::Repetition::OPTIONAL, parquet::Type::DOUBLE, parquet::ConvertedType::NONE); - auto list = parquet::schema::GroupNode::Make("list", parquet::Repetition::REPEATED, {element}); - fields.push_back(parquet::schema::GroupNode::Make(dsetname, parquet::Repetition::OPTIONAL, {list}, parquet::ConvertedType::LIST)); - } - std::shared_ptr schema = std::static_pointer_cast - (parquet::schema::GroupNode::Make("schema", parquet::Repetition::REQUIRED, fields)); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - int64_t i = 0; - int64_t numLeft = numelems; - auto segments = (int64_t*)chpl_segs; - int64_t valIdx = 0; // index into chpl_arr - int64_t segIdx = 0; // index into offsets - - if(dtype == ARROWINT64 || dtype == ARROWUINT64) { - while(numLeft > 0) { // write all local values to the file - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::Int64Writer* writer = - static_cast(rg_writer->NextColumn()); - int64_t count = 0; - while (numLeft > 0 && count < rowGroupSize) { // ensures rowGroupSize maintained - int64_t batchSize = segments[segIdx+1] - segments[segIdx]; - if (batchSize > 0) { - auto chpl_ptr = (int64_t*)chpl_arr; - int16_t* def_lvl = new int16_t[batchSize] { 3 }; // all values defined at the item level (3) - int16_t* rep_lvl = new int16_t[batchSize] { 0 }; - for (int64_t x = 0; x < batchSize; x++){ - // if the value is first in the segment rep_lvl = 0, otherwise 1 - rep_lvl[x] = (x == 0) ? 0 : 1; - def_lvl[x] = 3; - } - writer->WriteBatch(batchSize, def_lvl, rep_lvl, &chpl_ptr[valIdx]); - valIdx += batchSize; - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - batchSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - writer->WriteBatch(batchSize, &def_lvl, &rep_lvl, nullptr); - } - count++; - segIdx++; - numLeft--; - } - } - } - else if (dtype == ARROWBOOLEAN) { - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::BoolWriter* writer = - static_cast(rg_writer->NextColumn()); - int64_t count = 0; - while (numLeft > 0 && count < rowGroupSize) { - int64_t batchSize = segments[segIdx+1] - segments[segIdx]; - if (batchSize > 0) { - auto chpl_ptr = (bool*)chpl_arr; - // if the value is first in the segment rep_lvl = 0, otherwise 1 - // all values defined at the item level (3) - int16_t* def_lvl = new int16_t[batchSize] { 3 }; - int16_t* rep_lvl = new int16_t[batchSize] { 0 }; - for (int64_t x = 0; x < batchSize; x++){ - rep_lvl[x] = (x == 0) ? 0 : 1; - def_lvl[x] = 3; - } - writer->WriteBatch(batchSize, def_lvl, rep_lvl, &chpl_ptr[valIdx]); - valIdx += batchSize; - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - batchSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - writer->WriteBatch(batchSize, &def_lvl, &rep_lvl, nullptr); - } - count++; - segIdx++; - numLeft--; - } - } - } - else if (dtype == ARROWDOUBLE) { - while(numLeft > 0) { - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - parquet::DoubleWriter* writer = - static_cast(rg_writer->NextColumn()); - int64_t count = 0; - while (numLeft > 0 && count < rowGroupSize) { - int64_t batchSize = segments[segIdx+1] - segments[segIdx]; - if (batchSize > 0) { - auto chpl_ptr = (double*)chpl_arr; - // if the value is first in the segment rep_lvl = 0, otherwise 1 - // all values defined at the item level (3) - int16_t* def_lvl = new int16_t[batchSize] { 3 }; - int16_t* rep_lvl = new int16_t[batchSize] { 0 }; - for (int64_t x = 0; x < batchSize; x++){ - rep_lvl[x] = (x == 0) ? 0 : 1; - def_lvl[x] = 3; - } - writer->WriteBatch(batchSize, def_lvl, rep_lvl, &chpl_ptr[valIdx]); - valIdx += batchSize; - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - batchSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - writer->WriteBatch(batchSize, &def_lvl, &rep_lvl, nullptr); - } - count++; - segIdx++; - numLeft--; - } - } - } - else { - return ARROWERROR; - } - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -int cpp_writeMultiColToParquet(const char* filename, void* column_names, - void** ptr_arr, void** offset_arr, void* objTypes, void* datatypes, - void* segArr_sizes, int64_t colnum, int64_t numelems, int64_t rowGroupSize, - int64_t compression, char** errMsg) { - try { - // initialize the file to write to - using FileClass = ::arrow::io::FileOutputStream; - std::shared_ptr out_file; - ARROWRESULT_OK(FileClass::Open(filename), out_file); - - // Setup the parquet schema - std::shared_ptr schema = SetupSchema(column_names, objTypes, datatypes, colnum); - - parquet::WriterProperties::Builder builder; - // assign the proper compression - if(compression == SNAPPY_COMP) { - builder.compression(parquet::Compression::SNAPPY); - } else if (compression == GZIP_COMP) { - builder.compression(parquet::Compression::GZIP); - } else if (compression == BROTLI_COMP) { - builder.compression(parquet::Compression::BROTLI); - } else if (compression == ZSTD_COMP) { - builder.compression(parquet::Compression::ZSTD); - } else if (compression == LZ4_COMP) { - builder.compression(parquet::Compression::LZ4); - } - std::shared_ptr props = builder.build(); - - std::shared_ptr file_writer = - parquet::ParquetFileWriter::Open(out_file, schema, props); - - std::queue idxQueue_str; // queue used to track string byteIdx - std::queue idxQueue_segarray; // queue used to track index into the offsets - - auto dtypes_ptr = (int64_t*) datatypes; - auto objType_ptr = (int64_t*) objTypes; - auto saSizes_ptr = (int64_t*) segArr_sizes; - int64_t numLeft = numelems; // number of elements remaining to write (rows) - int64_t x = 0; // index to start writing batch from - while (numLeft > 0) { - // Append a RowGroup with a specific number of rows. - parquet::RowGroupWriter* rg_writer = file_writer->AppendRowGroup(); - int64_t batchSize = rowGroupSize; - if(numLeft < rowGroupSize) - batchSize = numLeft; - - // loop the columns and write the row groups - for(int64_t i = 0; i < colnum; i++){ - int64_t dtype = dtypes_ptr[i]; - if (dtype == ARROWINT64 || dtype == ARROWUINT64) { - auto data_ptr = (int64_t*)ptr_arr[i]; - parquet::Int64Writer* writer = - static_cast(rg_writer->NextColumn()); - - if (objType_ptr[i] == SEGARRAY) { - auto offset_ptr = (int64_t*)offset_arr[i]; - int64_t offIdx = 0; // index into offsets - - if (x > 0){ - offIdx = idxQueue_segarray.front(); - idxQueue_segarray.pop(); - } - - int64_t count = 0; - while (count < batchSize) { // ensures rowGroupSize maintained - int64_t segSize; - if (offIdx == (numelems - 1)) { - segSize = saSizes_ptr[i] - offset_ptr[offIdx]; - } - else { - segSize = offset_ptr[offIdx+1] - offset_ptr[offIdx]; - } - if (segSize > 0) { - int16_t* def_lvl = new int16_t[segSize] { 3 }; - int16_t* rep_lvl = new int16_t[segSize] { 0 }; - for (int64_t s = 0; s < segSize; s++){ - // if the value is first in the segment rep_lvl = 0, otherwise 1 - // all values defined at the item level (3) - rep_lvl[s] = (s == 0) ? 0 : 1; - def_lvl[s] = 3; - } - int64_t valIdx = offset_ptr[offIdx]; - writer->WriteBatch(segSize, def_lvl, rep_lvl, &data_ptr[valIdx]); - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - segSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - writer->WriteBatch(segSize, &def_lvl, &rep_lvl, nullptr); - } - offIdx++; - count++; - } - if (numLeft - count > 0) { - idxQueue_segarray.push(offIdx); - } - } else { - writer->WriteBatch(batchSize, nullptr, nullptr, &data_ptr[x]); - } - } else if(dtype == ARROWBOOLEAN) { - auto data_ptr = (bool*)ptr_arr[i]; - parquet::BoolWriter* writer = - static_cast(rg_writer->NextColumn()); - if (objType_ptr[i] == SEGARRAY) { - auto offset_ptr = (int64_t*)offset_arr[i]; - int64_t offIdx = 0; // index into offsets - - if (x > 0){ - offIdx = idxQueue_segarray.front(); - idxQueue_segarray.pop(); - } - - int64_t count = 0; - while (count < batchSize) { // ensures rowGroupSize maintained - int64_t segSize; - if (offIdx == numelems - 1) { - segSize = saSizes_ptr[i] - offset_ptr[offIdx]; - } - else { - segSize = offset_ptr[offIdx+1] - offset_ptr[offIdx]; - } - if (segSize > 0) { - int16_t* def_lvl = new int16_t[segSize] { 3 }; - int16_t* rep_lvl = new int16_t[segSize] { 0 }; - for (int64_t s = 0; s < segSize; s++){ - // if the value is first in the segment rep_lvl = 0, otherwise 1 - // all values defined at the item level (3) - rep_lvl[s] = (s == 0) ? 0 : 1; - def_lvl[s] = 3; - } - int64_t valIdx = offset_ptr[offIdx]; - writer->WriteBatch(segSize, def_lvl, rep_lvl, &data_ptr[valIdx]); - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - segSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl = 0; - writer->WriteBatch(segSize, &def_lvl, &rep_lvl, nullptr); - } - offIdx++; - count++; - } - if (numLeft - count > 0) { - idxQueue_segarray.push(offIdx); - } - } else { - writer->WriteBatch(batchSize, nullptr, nullptr, &data_ptr[x]); - } - } else if(dtype == ARROWDOUBLE) { - auto data_ptr = (double*)ptr_arr[i]; - parquet::DoubleWriter* writer = - static_cast(rg_writer->NextColumn()); - if (objType_ptr[i] == SEGARRAY) { - auto offset_ptr = (int64_t*)offset_arr[i]; - int64_t offIdx = 0; // index into offsets - - if (x > 0){ - offIdx = idxQueue_segarray.front(); - idxQueue_segarray.pop(); - } - - int64_t count = 0; - while (count < batchSize) { // ensures rowGroupSize maintained - int64_t segSize; - if (offIdx == numelems - 1) { - segSize = saSizes_ptr[i] - offset_ptr[offIdx]; - } - else { - segSize = offset_ptr[offIdx+1] - offset_ptr[offIdx]; - } - if (segSize > 0) { - int16_t* def_lvl = new int16_t[segSize] { 3 }; - int16_t* rep_lvl = new int16_t[segSize] { 0 }; - for (int64_t s = 0; s < segSize; s++){ - // if the value is first in the segment rep_lvl = 0, otherwise 1 - // all values defined at the item level (3) - rep_lvl[s] = (s == 0) ? 0 : 1; - def_lvl[s] = 3; - } - int64_t valIdx = offset_ptr[offIdx]; - writer->WriteBatch(segSize, def_lvl, rep_lvl, &data_ptr[valIdx]); - delete[] def_lvl; - delete[] rep_lvl; - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - segSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t def_lvl = 1; - int16_t rep_lvl =0; - writer->WriteBatch(segSize, &def_lvl, &rep_lvl, nullptr); - } - offIdx++; - count++; - } - if (numLeft - count > 0) { - idxQueue_segarray.push(offIdx); - } - } else { - writer->WriteBatch(batchSize, nullptr, nullptr, &data_ptr[x]); - } - } else if(dtype == ARROWSTRING) { - auto data_ptr = (uint8_t*)ptr_arr[i]; - parquet::ByteArrayWriter* ba_writer = - static_cast(rg_writer->NextColumn()); - if (objType_ptr[i] == SEGARRAY) { - auto offset_ptr = (int64_t*)offset_arr[i]; - int64_t byteIdx = 0; - int64_t offIdx = 0; // index into offsets - - // identify the starting byte index - if (x > 0){ - byteIdx = idxQueue_str.front(); - idxQueue_str.pop(); - - offIdx = idxQueue_segarray.front(); - idxQueue_segarray.pop(); - } - - int64_t count = 0; - while (count < batchSize) { // ensures rowGroupSize maintained - int64_t segSize; - if (offIdx == numelems - 1) { - segSize = saSizes_ptr[i] - offset_ptr[offIdx]; - } - else { - segSize = offset_ptr[offIdx+1] - offset_ptr[offIdx]; - } - if (segSize > 0) { - for (int64_t s=0; s(&data_ptr[byteIdx]); - int64_t nextIdx = byteIdx; - while (data_ptr[nextIdx] != 0x00){ - nextIdx++; - } - value.len = nextIdx - byteIdx; - ba_writer->WriteBatch(1, &def_lvl, &rep_lvl, &value); - byteIdx = nextIdx + 1; // increment to start of next word - } - } - else { - // empty segment denoted by null value that is not repeated (first of segment) defined at the list level (1) - segSize = 1; // even though segment is length=0, write null to hold the empty segment - int16_t* def_lvl = new int16_t[segSize] { 1 }; - int16_t* rep_lvl = new int16_t[segSize] { 0 }; - ba_writer->WriteBatch(segSize, def_lvl, rep_lvl, nullptr); - } - offIdx++; - count++; - } - if (numLeft - count > 0) { - idxQueue_str.push(byteIdx); - idxQueue_segarray.push(offIdx); - } - } - else { - int64_t count = 0; - int64_t byteIdx = 0; - - // identify the starting byte index - if (x > 0){ - byteIdx = idxQueue_str.front(); - idxQueue_str.pop(); - } - - while(count < batchSize) { - parquet::ByteArray value; - int16_t definition_level = 1; - value.ptr = reinterpret_cast(&data_ptr[byteIdx]); - int64_t nextIdx = byteIdx; - while (data_ptr[nextIdx] != 0x00){ - nextIdx++; - } - // subtract 1 since we have the null terminator - value.len = nextIdx - byteIdx; - ba_writer->WriteBatch(1, &definition_level, nullptr, &value); - count++; - byteIdx = nextIdx + 1; - } - if (numLeft - count > 0) { - idxQueue_str.push(byteIdx); - } - } - } else { - return ARROWERROR; - } - } - numLeft -= batchSize; - x += batchSize; - } - - file_writer->Close(); - ARROWSTATUS_OK(out_file->Close()); - - return 0; - } catch (const std::exception& e) { - *errMsg = strdup(e.what()); - return ARROWERROR; - } -} - -extern "C" { - int c_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - return cpp_writeColumnToParquet(filename, chpl_arr, colnum, dsetname, - numelems, rowGroupSize, dtype, compression, - errMsg); - } - int c_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - return cpp_writeStrColumnToParquet(filename, chpl_arr, chpl_offsets, - dsetname, numelems, rowGroupSize, dtype, compression, errMsg); - } - - int c_writeListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - return cpp_writeListColumnToParquet(filename, chpl_segs, chpl_arr, - dsetname, numelems, rowGroupSize, dtype, compression, errMsg); - } - - int c_writeStrListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_offsets, void* chpl_arr, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg) { - return cpp_writeStrListColumnToParquet(filename, chpl_segs, chpl_offsets, chpl_arr, - dsetname, numelems, rowGroupSize, dtype, compression, errMsg); - } - - int c_writeMultiColToParquet(const char* filename, void* column_names, - void** ptr_arr, void** offset_arr, void* objTypes, void* datatypes, - void* segArr_sizes, int64_t colnum, int64_t numelems, int64_t rowGroupSize, - int64_t compression, char** errMsg){ - return cpp_writeMultiColToParquet(filename, column_names, ptr_arr, offset_arr, objTypes, datatypes, segArr_sizes, colnum, numelems, rowGroupSize, compression, errMsg); - } -} diff --git a/src/parquet/WriteParquet.h b/src/parquet/WriteParquet.h deleted file mode 100644 index 7f42d462c3d..00000000000 --- a/src/parquet/WriteParquet.h +++ /dev/null @@ -1,57 +0,0 @@ -#include -#include - -// Wrap functions in C extern if compiling C++ object file -#ifdef __cplusplus -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -extern "C" { -#endif - - int cpp_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int c_writeColumnToParquet(const char* filename, void* chpl_arr, - int64_t colnum, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, char** errMsg); - - int c_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int cpp_writeStrColumnToParquet(const char* filename, void* chpl_arr, void* chpl_offsets, - const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - - int c_writeMultiColToParquet(const char* filename, void* column_names, - void** ptr_arr, void** offset_arr, void* objTypes, void* datatypes, - void* segArr_sizes, int64_t colnum, int64_t numelems, int64_t rowGroupSize, - int64_t compression, char** errMsg); - - int cpp_writeMultiColToParquet(const char* filename, void* column_names, - void** ptr_arr, void** offset_arr, void* objTypes, void* datatypes, - void* segArr_sizes, int64_t colnum, int64_t numelems, int64_t rowGroupSize, - int64_t compression, char** errMsg); - - int c_writeStrListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_offsets, - void* chpl_arr, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - int cpp_writeStrListColumnToParquet(const char* filename, void* chpl_segs, void* chpl_offsets, - void* chpl_arr, const char* dsetname, int64_t numelems, - int64_t rowGroupSize, int64_t dtype, int64_t compression, - char** errMsg); - -#ifdef __cplusplus -} -#endif diff --git a/tests/pandas/io_test.py b/tests/pandas/io_test.py index 241db82d504..bf04742e1b0 100644 --- a/tests/pandas/io_test.py +++ b/tests/pandas/io_test.py @@ -824,6 +824,12 @@ def test_multi_col_write(self, par_test_base_tmp, comp): rd_df = ak.DataFrame(rd_data) pd.testing.assert_frame_equal(akdf.to_pandas(), rd_df.to_pandas()) + # Arkouda Strings cannot represent null, so an empty string must be + # encoded as a defined zero-length value rather than a Parquet null. + for parquet_path in glob.glob(f"{tmp_dirname}/multi_col_parquet*"): + string_column = pq.read_table(parquet_path, columns=["c_9"])["c_9"] + assert string_column.null_count == 0 + # test save with index true akdf.to_parquet(f"{tmp_dirname}/idx_multi_col_parquet", index=True, compression=comp) rd_data = ak.read_parquet(f"{tmp_dirname}/idx_multi_col_parquet*") @@ -1050,13 +1056,7 @@ def test_against_standard_files(self): columns = ak.get_datasets(filename) assert columns == ans # Merely test that read succeeds, do not check output - if "delta_byte_array.parquet" not in filename: - ak.read_parquet(filename, datasets=columns) - else: - # Since delta encoding is not supported, the columns in - # this file should raise an error and not crash the server - with pytest.raises(RuntimeError): - ak.read_parquet(filename, datasets=columns) + ak.read_parquet(filename, datasets=columns) def test_null_handling_all(self, par_test_base_tmp): df = pd.DataFrame(