Skip to content

Instantly share code, notes, and snippets.

@suvayu
Last active December 13, 2018 21:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save suvayu/aa2d38cee82b97be76186ec00073fe10 to your computer and use it in GitHub Desktop.
Save suvayu/aa2d38cee82b97be76186ec00073fe10 to your computer and use it in GitHub Desktop.
Gandiva filtering MWE for SO (https://stackoverflow.com/q/53750405/289784)
CC := $(shell which gcc)
CXX := $(shell which g++)
PKGCONFIG := $(shell which pkg-config)
CXXSTD := -std=c++17
OPTS := -g -Wall -m64 $(CXXSTD)
CXXFLAGS := $(OPTS) -pthread -fPIC $(shell $(PKGCONFIG) --cflags gandiva)
LDFLAGS := $(shell $(PKGCONFIG) --libs gandiva)
%: %.cc
$(CXX) $(CXXFLAGS) $< -o $@ $(LDFLAGS)
#include <memory>
#include <iostream>
#include <arrow/memory_pool.h>
#include <arrow/array.h>
#include <arrow/builder.h>
#include <gandiva/arrow.h>
#include <gandiva/filter.h>
#include <gandiva/tree_expr_builder.h>
#include <arrow/status.h>
#define EXPECT_OK(expr) \
do { \
::arrow::Status s = (expr); \
} while (false)
using namespace arrow;
using namespace gandiva;
int main()
{
arrow::MemoryPool* pool_ = arrow::default_memory_pool();
int num_records = 5;
arrow::Int64Builder i64builder;
ArrayPtr array0;
EXPECT_OK(i64builder.AppendValues({1, 3, 2, 4, 5}));
EXPECT_OK(i64builder.Finish(&array0));
std::cout << array0->ToString() << std::endl;
// schema for input fields
auto field0 = field("f0", arrow::int64());
auto schema = arrow::schema({field0});
// even: f0 % 2 == 0
auto field0_node = TreeExprBuilder::MakeField(field0);
auto lit_2 = TreeExprBuilder::MakeLiteral(int64_t(2));
auto remainder = TreeExprBuilder::MakeFunction("mod", {field0_node, lit_2}, int64());
auto lit_0 = TreeExprBuilder::MakeLiteral(int64_t(0));
auto even = TreeExprBuilder::MakeFunction("equal", {remainder, lit_0}, boolean());
auto condition = TreeExprBuilder::MakeCondition(even);
// input record batch
auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0});
// filter
std::shared_ptr<Filter> filter;
EXPECT_OK(Filter::Make(schema, condition, &filter));
std::shared_ptr<SelectionVector> selected;
EXPECT_OK(SelectionVector::MakeInt16(num_records, pool_, &selected));
EXPECT_OK(filter->Evaluate(*in_batch, selected));
auto idx_arr = selected->ToArray();
std::cout << idx_arr->ToString() << "\n"
<< "length of idx_arr: " << idx_arr->length() << std::endl;
// std::cout << "array0[0]: " << array0->Value(0); // doesn't compile
// error: ‘using element_type = class arrow::Array’ {aka ‘class
// arrow::Array’} has no member named ‘Value’
// downcast it to the correct derived class
auto array0_cast = std::dynamic_pointer_cast<NumericArray<Int64Type>>(array0);
std::cout << "array0[0]: " << array0_cast->Value(0) << std::endl;
auto idx_arr_cast = std::dynamic_pointer_cast<NumericArray<Int16Type>>(idx_arr);
if (idx_arr_cast) {
std::cout << "idx_arr[0]: " << idx_arr_cast->Value(0) << std::endl;
} else {
std::cerr << "idx_arr_cast is a nullptr!" << std::endl;
}
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment