gaussian cubes infra

Author: cscheid

.gitignore

@@ -1,6 +1,7 @@
 *~
 .DS_Store
-
+CMakeCache.txt
+CMakeFiles
 build
 local_data
 bin

CMakeLists.txt

@@ -12,14 +12,13 @@ if(Boost_FOUND)
   include_directories(${Boost_INCLUDE_DIRS})
 endif()
 
-
 # include(cmake/FindArrow.cmake) findarrow doesn't work, so we'll do it by hand
+# let me say this will all my heart: fuck cmake
 
-set (ARROW_INCLUDE_DIR /Users/cscheid/repos/github/apache/arrow/cpp/src)
+set (ARROW_BASE_PATH /Users/cscheid/repos/github/apache/arrow/cpp)
+set (ARROW_INCLUDE_DIR ${ARROW_BASE_PATH}/src)
 # set (ARROW_LIBRARY_DIR /Users/cscheid/repos/github/apache/arrow/cpp/release/release)
-set (ARROW_LIBRARY_DIR /Users/cscheid/repos/github/apache/arrow/cpp/debug/debug)
-
-# let me say this will all my heart: fuck cmake
+set (ARROW_LIBRARY_DIR ${ARROW_BASE_PATH}/release/release)
 
 ################################################################################
 
@@ -97,15 +96,3 @@ target_include_directories(test_ztf_arrow PUBLIC ${ARROW_INCLUDE_DIR})
 target_link_directories(test_ztf_arrow PUBLIC ${ARROW_LIBRARY_DIR})
 
 target_link_libraries(test_ztf_arrow PUBLIC arrow)
-
-################################################################################
-
-add_executable(test_arrow_write
-  src/test_arrow_write
-  )
-
-target_include_directories(test_arrow_write PUBLIC ${ARROW_INCLUDE_DIR})
-
-target_link_directories(test_arrow_write PUBLIC ${ARROW_LIBRARY_DIR})
-
-target_link_libraries(test_arrow_write PUBLIC arrow)

client/js/agg.arrow

@@ -1,5 +1,40 @@
 import init, * as arrow from "./pkg/arrow_wasm.js";
 
+function vectorToArray(vec, type)
+{
+  if (type.name === 'floatingpoint' && type.precision === 'DOUBLE') {
+    return vec.asFloat64Vector().toArray();
+  } else if (type.name === 'floatingpoint' && type.precision === 'SINGLE') {
+    return vec.asFloat32Vector().toArray();
+  } else if (type.name === 'int' && type.bitWidth === 64 && type.isSigned) {
+    return vec.asInt64Vector().toArray();
+  } else if (type.name === 'int' && type.bitWidth === 64 && !type.isSigned) {
+    return vec.asUint64Vector().toArray();
+  } else if (type.name === 'int' && type.bitWidth === 32 && type.isSigned) {
+    return vec.asInt32Vector().toArray();
+  } else if (type.name === 'int' && type.bitWidth === 32 && !type.isSigned) {
+    return vec.asUint32Vector().toArray();
+  } else {
+    throw new Error(`Unrecognized type: ${JSON.stringify(type)}`);
+  }
+}
+
+function tableToDictOfArrays(table)
+{
+  let schema = table.schema.toJSON();
+  let result = {};
+  schema.fields.forEach((field, i) => {
+    let name = field.name;
+    let chunks = [];
+    for (let batch = 0; batch < table.numBatches; ++batch) {
+      chunks.push(vectorToArray(table.recordBatch(batch).column(i),
+                                schema.fields[i].type));
+    }
+    result[name] = chunks;
+  });
+  return result;
+}
+
 async function run() {
   await init();
 
@@ -9,8 +44,11 @@ async function run() {
   const data = await response.arrayBuffer();
   const contents = new Uint8Array(data);
 
-  const table = arrow.Table.from(contents);
+  let table = arrow.Table.from(contents);
   console.log(table.schema.toJSON());
+  table = tableToDictOfArrays(table);
+  
+  debugger;
 }
 
 run();

client/js/index.js

@@ -6,6 +6,8 @@
 
 #include <iostream>
 #include <algorithm>
+#include <sstream>
+#include <memory>
 
 #include "arrow_convenience.h"
 
@@ -42,6 +44,21 @@ shared_ptr<Table> make_table(
   return arrow::Table::Make(arrow::schema(fields), arrays);
 }
 
+void write_arrow(
+    shared_ptr<Table> table,
+    const string &path,
+    int64_t max_chunk_size)
+{
+  shared_ptr<io::FileOutputStream>
+      file = io::FileOutputStream::Open(path).ValueOrDie();
+
+  shared_ptr<ipc::RecordBatchWriter>
+      writer = ipc::MakeFileWriter(file, table->schema()).ValueOrDie();
+
+  OK_OR_DIE(writer->WriteTable(*table, max_chunk_size));
+  OK_OR_DIE(writer->Close());
+}
+
 /******************************************************************************/
 
 shared_ptr<Table> read_feather_table(const std::string& path)
@@ -248,3 +265,51 @@ sort_table(std::shared_ptr<Table> input,
       "sort_indices", { Datum(input) }, options).ValueOrDie().make_array();
   return permute_table(input, sort_permutation);
 }
+
+/******************************************************************************/
+
+shared_ptr<Table>
+make_gaussian_stats_table(
+    const vector<shared_ptr<ChunkedArray>> &cols)
+{
+  vector<NumericBuilder<DoubleType>>
+      builders((cols.size() + 1) * (cols.size() + 2) / 2);
+  int sz = (int) cols.size();
+  RowIterator(cols).for_each([&builders, &sz](RowIterator &rows) {
+    size_t ix = 0;
+    for (int i = -1; i < sz; ++i) {
+      double v1 = (i == -1) ? 1.0 : rows.cols_[i].value<DoubleType>();
+      for (int j = i; j < sz; ++j) {
+        double v2 = (j == -1) ? 1.0 : rows.cols_[j].value<DoubleType>();
+        OK_OR_DIE(builders[ix++].Append(v1 * v2));
+      }
+    }
+  });
+  unordered_map<string, shared_ptr<ChunkedArray>> result;
+  size_t ix = 0;
+  for (auto &builder: builders) {
+    stringstream ss;
+    ss << "s" << (ix++);
+    result[ss.str()] = shared_ptr<ChunkedArray>(new ChunkedArray({ builder.Finish().ValueOrDie() }));
+  }
+  return make_table(result);
+}
+
+shared_ptr<Table>
+make_gaussian_stats_table(
+    shared_ptr<Table> t,
+    const vector<string> &col_names)
+{
+  vector<shared_ptr<ChunkedArray>>
+      cols;
+  for (auto &name: col_names) {
+    cols.push_back(t->GetColumnByName(name));
+  }
+  return make_gaussian_stats_table(cols);
+}
+
+shared_ptr<Table>
+make_gaussian_stats_table(shared_ptr<Table> t)
+{
+  return make_gaussian_stats_table(t->columns());
+}

src/arrow_utils.cc

@@ -1,9 +1,13 @@
 #pragma once
 
 #include <arrow/compute/api_vector.h>
+#include <arrow/io/file.h>
+#include <arrow/io/api.h>
 #include <arrow/builder.h>
 #include <arrow/array.h>
 #include <arrow/table.h>
+#include <arrow/ipc/writer.h>
+
 #include <memory>
 #include <unordered_map>
 
@@ -29,6 +33,11 @@ void arrow_foreach(
     std::shared_ptr<arrow::ChunkedArray> column_2,
     T closure);
 
+void write_arrow(
+    std::shared_ptr<arrow::Table> table,
+    const std::string &path,
+    int64_t max_chunk_size = -1);
+
 /******************************************************************************/
 /// cbind: combines columns from two different tables.
 ///
@@ -255,6 +264,19 @@ struct ChunkedArrayIterator
 // Do not reuse RowIterators!
 struct RowIterator
 {
+  explicit RowIterator(std::shared_ptr<arrow::Table> table,
+                       const std::vector<std::string> &names,
+                       bool skip_null=true)
+  {
+    std::vector<std::shared_ptr<arrow::ChunkedArray>> cols;
+    for (auto &name: names) {
+      cols_.push_back(ChunkedArrayIterator(table->GetColumnByName(name)));
+    }
+    if (skip_null) {
+      ensure_not_null();
+    }
+  }
+  
   explicit RowIterator(const std::vector<ChunkedArrayIterator> &cols,
                        bool skip_null=true)
       : cols_(cols) {
@@ -364,6 +386,29 @@ template <typename ArrowType, typename T>
 std::shared_ptr<arrow::ChunkedArray>
 map_rows(RowIterator &rows, T closure);
 
+// same as above, but construct the RowIterator internally
+// given a table and some columns
+template <typename ArrowType, typename T>
+std::shared_ptr<arrow::ChunkedArray>
+map_rows(std::shared_ptr<arrow::Table> t,
+         const std::vector<std::string> &col_names,
+         T closure);
+
+/******************************************************************************/
+// create the table of sufficient statistics to fit gaussians.
+// this is (|cols|+1)(|cols|+2)/2
+//
+// the selected columns must all be DoubleType
+
+std::shared_ptr<arrow::Table>
+make_gaussian_stats_table(
+    std::shared_ptr<arrow::Table> t,
+    const std::vector<std::string> &cols);
+
+// This assumes all columns are to be selected.
+std::shared_ptr<arrow::Table>
+make_gaussian_stats_table(std::shared_ptr<arrow::Table> t);
+
 /******************************************************************************/
 
 #include "arrow_utils.hh"

src/arrow_utils.h

@@ -119,3 +119,20 @@ map_rows(RowIterator &rows, T closure)
   std::cerr << "Done!" << std::endl;
   return std::shared_ptr<arrow::ChunkedArray>(new arrow::ChunkedArray({ builder.Finish().ValueOrDie() }));
 }
+
+// same as above, but construct the RowIterator internally
+// given a table and some columns
+template <typename ArrowType, typename T>
+std::shared_ptr<arrow::ChunkedArray>
+map_rows(std::shared_ptr<arrow::Table> t,
+         const std::vector<std::string> &col_names,
+         T closure)
+{
+  std::vector<std::shared_ptr<arrow::ChunkedArray>> cols;
+  for (auto &name: col_names) {
+    cols.push_back(t->GetColumnByName(name));
+  }
+  RowIterator rows(cols);
+  return map_rows<ArrowType, T>(rows, closure);
+}
+