ArrowForeignStorage.cpp File Reference

#include "ArrowForeignStorage.h"
#include <arrow/api.h>
#include <arrow/csv/reader.h>
#include <arrow/io/file.h>
#include <arrow/util/decimal.h>
#include <tbb/parallel_for.h>
#include <tbb/task_group.h>
#include <array>
#include <future>
#include <vector>
#include "Catalog/DataframeTableDescriptor.h"
#include "DataMgr/ForeignStorage/ForeignStorageInterface.h"
#include "DataMgr/StringNoneEncoder.h"
#include "Logger/Logger.h"
#include "QueryEngine/ArrowResultSet.h"
#include "Shared/ArrowUtil.h"
#include "Shared/measure.h"

Include dependency graph for ArrowForeignStorage.cpp:

Go to the source code of this file.

Classes
struct	Frag
struct	ArrowFragment
class	ArrowForeignStorageBase
class	ArrowForeignStorage
class	ArrowCsvForeignStorage

Functions
void	convertBoolBitmapBufferWithNulls (int8_t *dst, const uint8_t *src, const uint8_t *bitmap, int64_t length, int8_t null_value)
void	convertBoolBitmapBufferWithoutNulls (int8_t *dst, const uint8_t *src, int64_t length)
template<typename V , std::enable_if_t<!std::is_same_v< V, bool > &&std::is_integral< V >::value, int > = 0>
V	get_null_value ()
template<typename V , std::enable_if_t< std::is_same_v< V, bool >, int > = 0>
int8_t	get_null_value ()
std::vector< Frag >	calculateFragmentsOffsets (const arrow::ChunkedArray &array, size_t maxFragRows)
static SQLTypeInfo	getOmnisciType (const arrow::DataType &type)
void	setArrowTable (std::string name, std::shared_ptr< arrow::Table > table)
void	releaseArrowTable (std::string name)
void	registerArrowForeignStorage (std::shared_ptr< ForeignStorageInterface > fsi)
static std::shared_ptr < arrow::DataType >	getArrowImportType (const SQLTypeInfo type)
void	registerArrowCsvForeignStorage (std::shared_ptr< ForeignStorageInterface > fsi)

Function Documentation

std::vector<Frag> calculateFragmentsOffsets	(	const arrow::ChunkedArray &	array,
		size_t	maxFragRows
	)

Definition at line 320 of file ArrowForeignStorage.cpp.

Referenced by ArrowForeignStorageBase::parseArrowTable().

                                                                {
  std::vector<Frag> fragments;
  size_t sz = 0;
  size_t offset = 0;
  fragments.push_back({0, 0, 0, 0});
  size_t num_chunks = (size_t)array.num_chunks();
  for (size_t i = 0; i < num_chunks;) {
    auto& chunk = *array.chunk(i);
    auto& frag = *fragments.rbegin();
    if (maxFragRows - sz > chunk.length() - offset) {
      sz += chunk.length() - offset;
      if (i == num_chunks - 1) {
        fragments.rbegin()->last_chunk = num_chunks - 1;
        fragments.rbegin()->last_chunk_size =
            array.chunk((int)num_chunks - 1)->length() - offset;
      }
      offset = 0;
      i++;
    } else {
      frag.last_chunk = i;
      frag.last_chunk_size = maxFragRows - sz;
      offset += maxFragRows - sz;
      sz = 0;
      fragments.push_back({i, offset, 0, 0});
    }
  }
  if (fragments.rbegin()->first_chunk == fragments.rbegin()->first_chunk &&
      fragments.rbegin()->last_chunk_size == 0) {
    // remove empty fragment at the end if any
    fragments.pop_back();
  }
  return fragments;
}

Here is the caller graph for this function:

void convertBoolBitmapBufferWithNulls	(	int8_t *	dst,
		const uint8_t *	src,
		const uint8_t *	bitmap,
		int64_t	length,
		int8_t	null_value
	)

Definition at line 142 of file ArrowForeignStorage.cpp.

References run_benchmark_import::dest, and is_null().

Referenced by ArrowForeignStorageBase::replaceNullValuesImpl().

                                                         {
  for (int64_t bitmap_idx = 0; bitmap_idx < length / 8; ++bitmap_idx) {
    auto source = src[bitmap_idx];
    auto dest = dst + bitmap_idx * 8;
    auto inversed_bitmap = ~bitmap[bitmap_idx];
    for (int8_t bitmap_offset = 0; bitmap_offset < 8; ++bitmap_offset) {
      auto is_null = (inversed_bitmap >> bitmap_offset) & 1;
      auto val = (source >> bitmap_offset) & 1;
      dest[bitmap_offset] = is_null ? null_value : val;
    }
  }

  for (int64_t j = (length / 8) * 8; j < length; ++j) {
    auto is_null = (~bitmap[length / 8] >> (j % 8)) & 1;
    auto val = (src[length / 8] >> (j % 8)) & 1;
    dst[j] = is_null ? null_value : val;
  }
}

Here is the call graph for this function:

Here is the caller graph for this function:

void convertBoolBitmapBufferWithoutNulls	(	int8_t *	dst,
		const uint8_t *	src,
		int64_t	length
	)

Definition at line 165 of file ArrowForeignStorage.cpp.

References run_benchmark_import::dest.

Referenced by ArrowForeignStorageBase::replaceNullValuesImpl().

                                                         {
  for (int64_t bitmap_idx = 0; bitmap_idx < length / 8; ++bitmap_idx) {
    auto source = src[bitmap_idx];
    auto dest = dst + bitmap_idx * 8;
    for (int8_t bitmap_offset = 0; bitmap_offset < 8; ++bitmap_offset) {
      dest[bitmap_offset] = (source >> bitmap_offset) & 1;
    }
  }

  for (int64_t j = (length / 8) * 8; j < length; ++j) {
    dst[j] = (src[length / 8] >> (j % 8)) & 1;
  }
}

Here is the caller graph for this function:

template<typename V , std::enable_if_t<!std::is_same_v< V, bool > &&std::is_integral< V >::value, int > = 0>

V get_null_value ( )

inline

Definition at line 184 of file ArrowForeignStorage.cpp.

Referenced by initializeFlatBuffer().

                          {
  return inline_int_null_value<V>();
}

Here is the caller graph for this function:

template<typename V , std::enable_if_t< std::is_same_v< V, bool >, int > = 0>

int8_t get_null_value ( )

inline

Definition at line 189 of file ArrowForeignStorage.cpp.

                               {
  return inline_int_null_value<int8_t>();
}

static std::shared_ptr<arrow::DataType> getArrowImportType ( const SQLTypeInfo  type )

static

Definition at line 924 of file ArrowForeignStorage.cpp.

References CHECK, SQLTypeInfo::get_precision(), SQLTypeInfo::get_scale(), SQLTypeInfo::get_size(), SQLTypeInfo::get_type(), SQLTypeInfo::get_type_name(), IS_INTEGER, kARRAY, kBOOLEAN, kCHAR, kDATE, kDECIMAL, kDOUBLE, kFLOAT, kINTERVAL_DAY_TIME, kINTERVAL_YEAR_MONTH, kNUMERIC, kTEXT, kTIME, kTIMESTAMP, kVARCHAR, and to_string().

Referenced by ArrowCsvForeignStorage::registerTable().

                                                                               {
  using namespace arrow;
  auto ktype = type.get_type();
  if (IS_INTEGER(ktype)) {
    switch (type.get_size()) {
      case 1:
        return int8();
      case 2:
        return int16();
      case 4:
        return int32();
      case 8:
        return int64();
      default:
        CHECK(false);
    }
  }
  switch (ktype) {
    case kBOOLEAN:
      return arrow::boolean();
    case kFLOAT:
      return float32();
    case kDOUBLE:
      return float64();
    case kCHAR:
    case kVARCHAR:
    case kTEXT:
      return utf8();
    case kDECIMAL:
    case kNUMERIC:
      return decimal(type.get_precision(), type.get_scale());
    case kTIME:
      return time32(TimeUnit::SECOND);
    case kDATE:
#ifdef HAVE_CUDA
      return arrow::date64();
#else
      return arrow::date32();
#endif
    case kTIMESTAMP:
      switch (type.get_precision()) {
        case 0:
          return timestamp(TimeUnit::SECOND);
        case 3:
          return timestamp(TimeUnit::MILLI);
        case 6:
          return timestamp(TimeUnit::MICRO);
        case 9:
          return timestamp(TimeUnit::NANO);
        default:
          throw std::runtime_error("Unsupported timestamp precision for Arrow: " +
                                   std::to_string(type.get_precision()));
      }
    case kARRAY:
    case kINTERVAL_DAY_TIME:
    case kINTERVAL_YEAR_MONTH:
    default:
      throw std::runtime_error(type.get_type_name() + " is not supported in Arrow.");
  }
  return nullptr;
}

Here is the call graph for this function:

Here is the caller graph for this function:

static SQLTypeInfo getOmnisciType ( const arrow::DataType &  type )

static

Definition at line 793 of file ArrowForeignStorage.cpp.

References kBIGINT, kBOOLEAN, kDATE, kDECIMAL, kDOUBLE, kENCODING_DICT, kFLOAT, kINT, kSMALLINT, kTEXT, kTIME, kTIMESTAMP, kTINYINT, heavydb.dtypes::STRING, and run_benchmark_import::type.

Referenced by ArrowForeignStorage::prepareTable().

                                                             {
  using namespace arrow;
  switch (type.id()) {
    case Type::INT8:
      return SQLTypeInfo(kTINYINT, false);
    case Type::INT16:
      return SQLTypeInfo(kSMALLINT, false);
    case Type::INT32:
      return SQLTypeInfo(kINT, false);
    case Type::INT64:
      return SQLTypeInfo(kBIGINT, false);
    case Type::BOOL:
      return SQLTypeInfo(kBOOLEAN, false);
    case Type::FLOAT:
      return SQLTypeInfo(kFLOAT, false);
    case Type::DATE32:
    case Type::DATE64:
      return SQLTypeInfo(kDATE, false);
    case Type::DOUBLE:
      return SQLTypeInfo(kDOUBLE, false);
      // uncomment when arrow 2.0 will be released and modin support for dictionary types
      // in read_csv would be implemented

      // case Type::DICTIONARY: {
      //   auto type = SQLTypeInfo(kTEXT, false, kENCODING_DICT);
      //   // this is needed because createTable forces type.size to be equal to
      //   // comp_param / 8, no matter what type.size you set here
      //   type.set_comp_param(sizeof(uint32_t) * 8);
      //   return type;
      // }
      // case Type::STRING:
      //   return SQLTypeInfo(kTEXT, false, kENCODING_NONE);

    case Type::STRING: {
      auto type = SQLTypeInfo(kTEXT, false, kENCODING_DICT);
      // this is needed because createTable forces type.size to be equal to
      // comp_param / 8, no matter what type.size you set here
      type.set_comp_param(sizeof(uint32_t) * 8);
      return type;
    }
    case Type::DECIMAL: {
      const auto& decimal_type = static_cast<const arrow::DecimalType&>(type);
      return SQLTypeInfo(kDECIMAL, decimal_type.precision(), decimal_type.scale(), false);
    }
    case Type::TIME32:
      return SQLTypeInfo(kTIME, false);
    case Type::TIMESTAMP:
      switch (static_cast<const arrow::TimestampType&>(type).unit()) {
        case TimeUnit::SECOND:
          return SQLTypeInfo(kTIMESTAMP, 0, 0);
        case TimeUnit::MILLI:
          return SQLTypeInfo(kTIMESTAMP, 3, 0);
        case TimeUnit::MICRO:
          return SQLTypeInfo(kTIMESTAMP, 6, 0);
        case TimeUnit::NANO:
          return SQLTypeInfo(kTIMESTAMP, 9, 0);
      }
    default:
      throw std::runtime_error(type.ToString() + " is not yet supported.");
  }
}

Here is the caller graph for this function:

void registerArrowCsvForeignStorage

(

std::shared_ptr< ForeignStorageInterface >

fsi

)

Definition at line 1063 of file ArrowForeignStorage.cpp.

Referenced by PersistentStorageMgr::PersistentStorageMgr().

                                                                                {
  fsi->registerPersistentStorageInterface(std::make_unique<ArrowCsvForeignStorage>());
}

Here is the caller graph for this function:

void registerArrowForeignStorage

(

std::shared_ptr< ForeignStorageInterface >

fsi

)

Definition at line 893 of file ArrowForeignStorage.cpp.

Referenced by PersistentStorageMgr::PersistentStorageMgr().

                                                                             {
  fsi->registerPersistentStorageInterface(std::make_unique<ArrowForeignStorage>());
}

Here is the caller graph for this function:

void releaseArrowTable

(

std::string

name

)

Definition at line 889 of file ArrowForeignStorage.cpp.

References ArrowForeignStorage::tables.

Referenced by EmbeddedDatabase::DBEngineImpl::importArrowTable().

                                       {
  ArrowForeignStorage::tables.erase(name);
}

Here is the caller graph for this function:

void setArrowTable	(	std::string	name,
		std::shared_ptr< arrow::Table >	table
	)

Definition at line 885 of file ArrowForeignStorage.cpp.

References setup::name, and ArrowForeignStorage::tables.

Referenced by EmbeddedDatabase::DBEngineImpl::importArrowTable().

                                                                    {
  ArrowForeignStorage::tables[name] = table;
}

Here is the caller graph for this function: