Complete dictionary support

.github/workflows/r.yml

@@ -92,21 +92,20 @@ jobs:
         continue-on-error: true
         run: archery docker push ubuntu-r
 
-  rstudio:
-    name: "rstudio/r-base:${{ matrix.r_version }}-${{ matrix.r_image }}"
+  bundled:
+    name: "${{ matrix.config.org }}/${{ matrix.config.image }}:${{ matrix.config.tag }}"
     runs-on: ubuntu-latest
     if: ${{ !contains(github.event.pull_request.title, 'WIP') }}
     strategy:
       fail-fast: false
       matrix:
-        # See https://hub.docker.com/r/rstudio/r-base
-        r_version: ["4.0"]
-        r_image:
-          - centos7
+        config:
+          - {org: 'rstudio', image: 'r-base', tag: '4.0-centos7'}
+          - {org: 'rhub', image: 'debian-gcc-devel', tag: 'latest'}
     env:
-      R_ORG: rstudio
-      R_IMAGE: r-base
-      R_TAG: ${{ matrix.r_version }}-${{ matrix.r_image }}
+      R_ORG: ${{ matrix.config.org }}
+      R_IMAGE: ${{ matrix.config.image }}
+      R_TAG: ${{ matrix.config.tag }}
     steps:
       - name: Checkout Arrow
         uses: actions/checkout@v2
@@ -120,8 +119,8 @@ jobs:
         uses: actions/cache@v1
         with:
           path: .docker
-          key: ${{ matrix.r_image }}-r-${{ hashFiles('cpp/**') }}
-          restore-keys: ${{ matrix.r_image }}-r-
+          key: ${{ matrix.config.image }}-r-${{ hashFiles('cpp/**') }}
+          restore-keys: ${{ matrix.config.image }}-r-
       - name: Setup Python
         uses: actions/setup-python@v1
         with:

cpp/cmake_modules/SetupCxxFlags.cmake

@@ -18,6 +18,7 @@
 # Check if the target architecture and compiler supports some special
 # instruction sets that would boost performance.
 include(CheckCXXCompilerFlag)
+include(CheckCXXSourceCompiles)
 # Get cpu architecture
 
 message(STATUS "System processor: ${CMAKE_SYSTEM_PROCESSOR}")
@@ -60,17 +61,36 @@ if(ARROW_CPU_FLAG STREQUAL "x86")
     # https://gcc.gnu.org/bugzilla/show_bug.cgi?id=65782
     message(STATUS "Disable AVX512 support on MINGW for now")
   else()
-    check_cxx_compiler_flag(${ARROW_AVX512_FLAG} CXX_SUPPORTS_AVX512)
+    # Check for AVX512 support in the compiler.
+    set(OLD_CMAKE_REQURED_FLAGS ${CMAKE_REQUIRED_FLAGS})
+    set(CMAKE_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS} ${ARROW_AVX512_FLAG}")
+    check_cxx_source_compiles("
+      #ifdef _MSC_VER
+      #include <intrin.h>
+      #else
+      #include <immintrin.h>
+      #endif
+
+      int main() {
+        __m512i mask = _mm512_set1_epi32(0x1);
+        char out[32];
+        _mm512_storeu_si512(out, mask);
+        return 0;
+      }" CXX_SUPPORTS_AVX512)
+    set(CMAKE_REQUIRED_FLAGS ${OLD_CMAKE_REQURED_FLAGS})
   endif()
-  # Runtime SIMD level it can get from compiler
+  # Runtime SIMD level it can get from compiler and ARROW_RUNTIME_SIMD_LEVEL
   if(CXX_SUPPORTS_SSE4_2
      AND ARROW_RUNTIME_SIMD_LEVEL MATCHES "^(SSE4_2|AVX2|AVX512|MAX)$")
+    set(ARROW_HAVE_RUNTIME_SSE4_2 ON)
     add_definitions(-DARROW_HAVE_RUNTIME_SSE4_2)
   endif()
   if(CXX_SUPPORTS_AVX2 AND ARROW_RUNTIME_SIMD_LEVEL MATCHES "^(AVX2|AVX512|MAX)$")
+    set(ARROW_HAVE_RUNTIME_AVX2 ON)
     add_definitions(-DARROW_HAVE_RUNTIME_AVX2 -DARROW_HAVE_RUNTIME_BMI2)
   endif()
   if(CXX_SUPPORTS_AVX512 AND ARROW_RUNTIME_SIMD_LEVEL MATCHES "^(AVX512|MAX)$")
+    set(ARROW_HAVE_RUNTIME_AVX512 ON)
     add_definitions(-DARROW_HAVE_RUNTIME_AVX512 -DARROW_HAVE_RUNTIME_BMI2)
   endif()
 elseif(ARROW_CPU_FLAG STREQUAL "ppc")

cpp/src/arrow/CMakeLists.txt

@@ -214,13 +214,13 @@ set(ARROW_SRCS
     vendored/double-conversion/diy-fp.cc
     vendored/double-conversion/strtod.cc)
 
-if(CXX_SUPPORTS_AVX2)
+if(ARROW_HAVE_RUNTIME_AVX2)
   list(APPEND ARROW_SRCS util/bpacking_avx2.cc)
   set_source_files_properties(util/bpacking_avx2.cc PROPERTIES SKIP_PRECOMPILE_HEADERS ON)
   set_source_files_properties(util/bpacking_avx2.cc PROPERTIES COMPILE_FLAGS
                               ${ARROW_AVX2_FLAG})
 endif()
-if(CXX_SUPPORTS_AVX512)
+if(ARROW_HAVE_RUNTIME_AVX512)
   list(APPEND ARROW_SRCS util/bpacking_avx512.cc)
   set_source_files_properties(util/bpacking_avx512.cc PROPERTIES SKIP_PRECOMPILE_HEADERS
                               ON)
@@ -387,14 +387,14 @@ if(ARROW_COMPUTE)
               compute/kernels/vector_selection.cc
               compute/kernels/vector_sort.cc)
 
-  if(CXX_SUPPORTS_AVX2)
+  if(ARROW_HAVE_RUNTIME_AVX2)
     list(APPEND ARROW_SRCS compute/kernels/aggregate_basic_avx2.cc)
     set_source_files_properties(compute/kernels/aggregate_basic_avx2.cc PROPERTIES
                                 SKIP_PRECOMPILE_HEADERS ON)
     set_source_files_properties(compute/kernels/aggregate_basic_avx2.cc PROPERTIES
                                 COMPILE_FLAGS ${ARROW_AVX2_FLAG})
   endif()
-  if(CXX_SUPPORTS_AVX512)
+  if(ARROW_HAVE_RUNTIME_AVX512)
     list(APPEND ARROW_SRCS compute/kernels/aggregate_basic_avx512.cc)
     set_source_files_properties(compute/kernels/aggregate_basic_avx512.cc PROPERTIES
                                 SKIP_PRECOMPILE_HEADERS ON)

cpp/src/arrow/compute/kernels/aggregate_test.cc

@@ -953,14 +953,13 @@ class TestPrimitiveVarStdKernel : public ::testing::Test {
   using ScalarType = typename TypeTraits<DoubleType>::ScalarType;
 
   void AssertVarStdIs(const Array& array, const VarianceOptions& options,
-                      double expected_var, double diff = 0) {
-    AssertVarStdIsInternal(array, options, expected_var, diff);
+                      double expected_var) {
+    AssertVarStdIsInternal(array, options, expected_var);
   }
 
   void AssertVarStdIs(const std::shared_ptr<ChunkedArray>& array,
-                      const VarianceOptions& options, double expected_var,
-                      double diff = 0) {
-    AssertVarStdIsInternal(array, options, expected_var, diff);
+                      const VarianceOptions& options, double expected_var) {
+    AssertVarStdIsInternal(array, options, expected_var);
   }
 
   void AssertVarStdIs(const std::string& json, const VarianceOptions& options,
@@ -999,18 +998,14 @@ class TestPrimitiveVarStdKernel : public ::testing::Test {
 
  private:
   void AssertVarStdIsInternal(const Datum& array, const VarianceOptions& options,
-                              double expected_var, double diff = 0) {
+                              double expected_var) {
     ASSERT_OK_AND_ASSIGN(Datum out_var, Variance(array, options));
     ASSERT_OK_AND_ASSIGN(Datum out_std, Stddev(array, options));
     auto var = checked_cast<const ScalarType*>(out_var.scalar().get());
     auto std = checked_cast<const ScalarType*>(out_std.scalar().get());
     ASSERT_TRUE(var->is_valid && std->is_valid);
     ASSERT_DOUBLE_EQ(std->value * std->value, var->value);
-    if (diff == 0) {
-      ASSERT_DOUBLE_EQ(var->value, expected_var);  // < 4ULP
-    } else {
-      ASSERT_NEAR(var->value, expected_var, diff);
-    }
+    ASSERT_DOUBLE_EQ(var->value, expected_var);  // < 4ULP
   }
 
   void AssertVarStdIsInvalidInternal(const Datum& array, const VarianceOptions& options) {
@@ -1070,22 +1065,39 @@ TEST_F(TestVarStdKernelStability, Basics) {
   VarianceOptions options{1};  // ddof = 1
   this->AssertVarStdIs("[100000004, 100000007, 100000013, 100000016]", options, 30.0);
   this->AssertVarStdIs("[1000000004, 1000000007, 1000000013, 1000000016]", options, 30.0);
+
+#ifndef __MINGW32__  // MinGW has precision issues
+  // This test is to make sure our variance combining method is stable.
+  // XXX: The reference value from numpy is actually wrong due to floating
+  // point limits. The correct result should equals variance(90, 0) = 4050.
+  std::vector<std::string> chunks = {"[40000008000000490]", "[40000008000000400]"};
+  this->AssertVarStdIs(chunks, options, 3904.0);
+#endif
+}
+
+// https://en.wikipedia.org/wiki/Kahan_summation_algorithm
+void KahanSum(double& sum, double& adjust, double addend) {
+  double y = addend - adjust;
+  double t = sum + y;
+  adjust = (t - sum) - y;
+  sum = t;
 }
 
-// Calculate reference variance with Welford's online algorithm
+// Calculate reference variance with Welford's online algorithm + Kahan summation
 // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
 std::pair<double, double> WelfordVar(const Array& array) {
   const auto& array_numeric = reinterpret_cast<const DoubleArray&>(array);
   const auto values = array_numeric.raw_values();
   internal::BitmapReader reader(array.null_bitmap_data(), array.offset(), array.length());
   double count = 0, mean = 0, m2 = 0;
+  double mean_adjust = 0, m2_adjust = 0;
   for (int64_t i = 0; i < array.length(); ++i) {
     if (reader.IsSet()) {
       ++count;
       double delta = values[i] - mean;
-      mean += delta / count;
+      KahanSum(mean, mean_adjust, delta / count);
       double delta2 = values[i] - mean;
-      m2 += delta * delta2;
+      KahanSum(m2, m2_adjust, delta * delta2);
     }
     reader.Next();
   }
@@ -1116,8 +1128,8 @@ TEST_F(TestVarStdKernelRandom, Basics) {
   double var_population, var_sample;
   std::tie(var_population, var_sample) = WelfordVar(*(array->Slice(0, total_size)));
 
-  this->AssertVarStdIs(chunked, VarianceOptions{0}, var_population, 0.0001);
-  this->AssertVarStdIs(chunked, VarianceOptions{1}, var_sample, 0.0001);
+  this->AssertVarStdIs(chunked, VarianceOptions{0}, var_population);
+  this->AssertVarStdIs(chunked, VarianceOptions{1}, var_sample);
 }
 
 }  // namespace compute

cpp/src/arrow/compute/kernels/aggregate_var_std.cc

@@ -53,32 +53,33 @@ struct VarStdState {
         []() {});
 
     this->count = count;
-    this->sum = sum;
+    this->mean = mean;
     this->m2 = m2;
   }
 
-  // Combine `m2` from two chunks
-  // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
+  // Combine `m2` from two chunks (m2 = n*s2)
+  // https://www.emathzone.com/tutorials/basic-statistics/combined-variance.html
   void MergeFrom(const ThisType& state) {
     if (state.count == 0) {
       return;
     }
     if (this->count == 0) {
       this->count = state.count;
-      this->sum = state.sum;
+      this->mean = state.mean;
       this->m2 = state.m2;
       return;
     }
-    double delta = this->sum / this->count - state.sum / state.count;
-    this->m2 += state.m2 +
-                delta * delta * this->count * state.count / (this->count + state.count);
+    double mean = (this->mean * this->count + state.mean * state.count) /
+                  (this->count + state.count);
+    this->m2 += state.m2 + this->count * (this->mean - mean) * (this->mean - mean) +
+                state.count * (state.mean - mean) * (state.mean - mean);
     this->count += state.count;
-    this->sum += state.sum;
+    this->mean = mean;
   }
 
   int64_t count = 0;
-  double sum = 0;
-  double m2 = 0;  // sum((X-mean)^2)
+  double mean = 0;
+  double m2 = 0;  // m2 = count*s2 = sum((X-mean)^2)
 };
 
 enum class VarOrStd : bool { Var, Std };

cpp/src/arrow/dataset/partition_test.cc

@@ -156,6 +156,9 @@ TEST_F(TestPartitioning, DiscoverSchema) {
   // fall back to string if any segment for field alpha is not parseable as int
   AssertInspect({"/0/1", "/hello/1"}, {Str("alpha"), Int("beta")});
 
+  // If there are too many digits fall back to string
+  AssertInspect({"/3760212050/1"}, {Str("alpha"), Int("beta")});
+
   // missing segment for beta doesn't cause an error or fallback
   AssertInspect({"/0/1", "/hello"}, {Str("alpha"), Int("beta")});
 }
@@ -168,6 +171,9 @@ TEST_F(TestPartitioning, DictionaryInference) {
   // type is still int32 if possible
   AssertInspect({"/0/1"}, {DictInt("alpha"), DictInt("beta")});
 
+  // If there are too many digits fall back to string
+  AssertInspect({"/3760212050/1"}, {DictStr("alpha"), DictInt("beta")});
+
   // successful dictionary inference
   AssertInspect({"/a/0"}, {DictStr("alpha"), DictInt("beta")});
   AssertInspect({"/a/0", "/a/1"}, {DictStr("alpha"), DictInt("beta")});
@@ -256,6 +262,9 @@ TEST_F(TestPartitioning, DiscoverHiveSchema) {
   // (...so ensure your partitions are ordered the same for all paths)
   AssertInspect({"/alpha=0/beta=1", "/beta=2/alpha=3"}, {Int("alpha"), Int("beta")});
 
+  // If there are too many digits fall back to string
+  AssertInspect({"/alpha=3760212050"}, {Str("alpha")});
+
   // missing path segments will not cause an error
   AssertInspect({"/alpha=0/beta=1", "/beta=2/alpha=3", "/gamma=what"},
                 {Int("alpha"), Int("beta"), Str("gamma")});
@@ -269,6 +278,9 @@ TEST_F(TestPartitioning, HiveDictionaryInference) {
   // type is still int32 if possible
   AssertInspect({"/alpha=0/beta=1"}, {DictInt("alpha"), DictInt("beta")});
 
+  // If there are too many digits fall back to string
+  AssertInspect({"/alpha=3760212050"}, {DictStr("alpha")});
+
   // successful dictionary inference
   AssertInspect({"/alpha=a/beta=0"}, {DictStr("alpha"), DictInt("beta")});
   AssertInspect({"/alpha=a/beta=0", "/alpha=a/1"}, {DictStr("alpha"), DictInt("beta")});

cpp/src/arrow/ipc/metadata_internal.cc

@@ -427,8 +427,7 @@ static Status GetDictionaryEncoding(FBB& fbb, const std::shared_ptr<Field>& fiel
                                     const DictionaryType& type, int64_t dictionary_id,
                                     DictionaryOffset* out) {
   // We assume that the dictionary index type (as an integer) has already been
-  // validated elsewhere, and can safely assume we are dealing with signed
-  // integers
+  // validated elsewhere, and can safely assume we are dealing with integers
   const auto& index_type = checked_cast<const IntegerType&>(*type.index_type());
 
   auto index_type_offset =

cpp/src/arrow/util/utf8.h

@@ -27,6 +27,7 @@
 #include "arrow/util/macros.h"
 #include "arrow/util/simd.h"
 #include "arrow/util/string_view.h"
+#include "arrow/util/ubsan.h"
 #include "arrow/util/visibility.h"
 
 namespace arrow {
@@ -87,8 +88,9 @@ ARROW_EXPORT void InitializeUTF8();
 
 inline bool ValidateUTF8(const uint8_t* data, int64_t size) {
   static constexpr uint64_t high_bits_64 = 0x8080808080808080ULL;
-  // For some reason, defining this variable outside the loop helps clang
-  uint64_t mask;
+  static constexpr uint32_t high_bits_32 = 0x80808080UL;
+  static constexpr uint16_t high_bits_16 = 0x8080U;
+  static constexpr uint8_t high_bits_8 = 0x80U;
 
 #ifndef NDEBUG
   internal::CheckUTF8Initialized();
@@ -98,8 +100,8 @@ inline bool ValidateUTF8(const uint8_t* data, int64_t size) {
     // XXX This is doing an unaligned access.  Contemporary architectures
     // (x86-64, AArch64, PPC64) support it natively and often have good
     // performance nevertheless.
-    memcpy(&mask, data, 8);
-    if (ARROW_PREDICT_TRUE((mask & high_bits_64) == 0)) {
+    uint64_t mask64 = SafeLoadAs<uint64_t>(data);
+    if (ARROW_PREDICT_TRUE((mask64 & high_bits_64) == 0)) {
       // 8 bytes of pure ASCII, move forward
       size -= 8;
       data += 8;
@@ -154,13 +156,50 @@ inline bool ValidateUTF8(const uint8_t* data, int64_t size) {
     return false;
   }
 
-  // Validate string tail one byte at a time
+  // Check if string tail is full ASCII (common case, fast)
+  if (size >= 4) {
+    uint32_t tail_mask = SafeLoadAs<uint32_t>(data + size - 4);
+    uint32_t head_mask = SafeLoadAs<uint32_t>(data);
+    if (ARROW_PREDICT_TRUE(((head_mask | tail_mask) & high_bits_32) == 0)) {
+      return true;
+    }
+  } else if (size >= 2) {
+    uint16_t tail_mask = SafeLoadAs<uint16_t>(data + size - 2);
+    uint16_t head_mask = SafeLoadAs<uint16_t>(data);
+    if (ARROW_PREDICT_TRUE(((head_mask | tail_mask) & high_bits_16) == 0)) {
+      return true;
+    }
+  } else if (size == 1) {
+    if (ARROW_PREDICT_TRUE((*data & high_bits_8) == 0)) {
+      return true;
+    }
+  } else {
+    /* size == 0 */
+    return true;
+  }
+
+  // Fall back to UTF8 validation of tail string.
   // Note the state table is designed so that, once in the reject state,
   // we remain in that state until the end.  So we needn't check for
   // rejection at each char (we don't gain much by short-circuiting here).
   uint16_t state = internal::kUTF8ValidateAccept;
-  while (size-- > 0) {
-    state = internal::ValidateOneUTF8Byte(*data++, state);
+  switch (size) {
+    case 7:
+      state = internal::ValidateOneUTF8Byte(data[size - 7], state);
+    case 6:
+      state = internal::ValidateOneUTF8Byte(data[size - 6], state);
+    case 5:
+      state = internal::ValidateOneUTF8Byte(data[size - 5], state);
+    case 4:
+      state = internal::ValidateOneUTF8Byte(data[size - 4], state);
+    case 3:
+      state = internal::ValidateOneUTF8Byte(data[size - 3], state);
+    case 2:
+      state = internal::ValidateOneUTF8Byte(data[size - 2], state);
+    case 1:
+      state = internal::ValidateOneUTF8Byte(data[size - 1], state);
+    default:
+      break;
   }
   return ARROW_PREDICT_TRUE(state == internal::kUTF8ValidateAccept);
 }