+chunked_array and tests

romainfrancois · romainfrancois · commit e02e24f76432 · 2018-09-24T20:15:02.000+02:00
diff --git a/r/NAMESPACE b/r/NAMESPACE
@@ -18,6 +18,7 @@ export(Type)
 export(array)
 export(as_tibble)
 export(boolean)
+export(chunked_array)
 export(date32)
 export(date64)
 export(decimal)
diff --git a/r/R/RcppExports.R b/r/R/RcppExports.R
@@ -117,6 +117,18 @@ ChunkedArray__as_vector <- function(chunked_array) {
     .Call(`_arrow_ChunkedArray__as_vector`, chunked_array)
 }
 
+ChunkArray__Slice1 <- function(chunked_array, offset) {
+    .Call(`_arrow_ChunkArray__Slice1`, chunked_array, offset)
+}
+
+ChunkArray__Slice2 <- function(chunked_array, offset, length) {
+    .Call(`_arrow_ChunkArray__Slice2`, chunked_array, offset, length)
+}
+
+ChunkedArray__Make <- function(chunks) {
+    .Call(`_arrow_ChunkedArray__Make`, chunks)
+}
+
 RecordBatch_to_dataframe <- function(batch) {
     .Call(`_arrow_RecordBatch_to_dataframe`, batch)
 }
diff --git a/r/R/array.R b/r/R/array.R
@@ -108,7 +108,14 @@ read_record_batch <- function(path){
     chunk = function(i) `arrow::Array`$new(ChunkedArray__chunk(self, i)),
     chunks = function() purrr::map(ChunkedArray__chunks(self), `arrow::Array`$new),
     type = function() `arrow::DataType`$dispatch(ChunkedArray__type(self)),
-    as_vector = function() ChunkedArray__as_vector(self)
+    as_vector = function() ChunkedArray__as_vector(self),
+    Slice = function(offset, length = NULL){
+      if (is.null(length)) {
+        `arrow::ChunkedArray`$new(ChunkArray__Slice1(self, offset))
+      } else {
+        `arrow::ChunkedArray`$new(ChunkArray__Slice2(self, offset, length))
+      }
+    }
   )
 )
 
@@ -161,4 +168,11 @@ read_arrow <- function(path){
   as_tibble(read_table(path))
 }
 
-
+#' create an arrow::Array from an R vector
+#'
+#' @param \dots Vectors to coerce
+#'
+#' @export
+chunked_array <- function(...){
+  `arrow::ChunkedArray`$new(ChunkedArray__Make(list(...)))
+}
diff --git a/r/man/chunked_array.Rd b/r/man/chunked_array.Rd
diff --git a/r/src/RcppExports.cpp b/r/src/RcppExports.cpp
@@ -331,6 +331,42 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// ChunkArray__Slice1
+std::shared_ptr<arrow::ChunkedArray> ChunkArray__Slice1(const std::shared_ptr<arrow::ChunkedArray>& chunked_array, int offset);
+RcppExport SEXP _arrow_ChunkArray__Slice1(SEXP chunked_arraySEXP, SEXP offsetSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const std::shared_ptr<arrow::ChunkedArray>& >::type chunked_array(chunked_arraySEXP);
+    Rcpp::traits::input_parameter< int >::type offset(offsetSEXP);
+    rcpp_result_gen = Rcpp::wrap(ChunkArray__Slice1(chunked_array, offset));
+    return rcpp_result_gen;
+END_RCPP
+}
+// ChunkArray__Slice2
+std::shared_ptr<arrow::ChunkedArray> ChunkArray__Slice2(const std::shared_ptr<arrow::ChunkedArray>& chunked_array, int offset, int length);
+RcppExport SEXP _arrow_ChunkArray__Slice2(SEXP chunked_arraySEXP, SEXP offsetSEXP, SEXP lengthSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const std::shared_ptr<arrow::ChunkedArray>& >::type chunked_array(chunked_arraySEXP);
+    Rcpp::traits::input_parameter< int >::type offset(offsetSEXP);
+    Rcpp::traits::input_parameter< int >::type length(lengthSEXP);
+    rcpp_result_gen = Rcpp::wrap(ChunkArray__Slice2(chunked_array, offset, length));
+    return rcpp_result_gen;
+END_RCPP
+}
+// ChunkedArray__Make
+std::shared_ptr<arrow::ChunkedArray> ChunkedArray__Make(List chunks);
+RcppExport SEXP _arrow_ChunkedArray__Make(SEXP chunksSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< List >::type chunks(chunksSEXP);
+    rcpp_result_gen = Rcpp::wrap(ChunkedArray__Make(chunks));
+    return rcpp_result_gen;
+END_RCPP
+}
 // RecordBatch_to_dataframe
 List RecordBatch_to_dataframe(const std::shared_ptr<arrow::RecordBatch>& batch);
 RcppExport SEXP _arrow_RecordBatch_to_dataframe(SEXP batchSEXP) {
@@ -1090,6 +1126,9 @@ static const R_CallMethodDef CallEntries[] = {
     {"_arrow_ChunkedArray__chunks", (DL_FUNC) &_arrow_ChunkedArray__chunks, 1},
     {"_arrow_ChunkedArray__type", (DL_FUNC) &_arrow_ChunkedArray__type, 1},
     {"_arrow_ChunkedArray__as_vector", (DL_FUNC) &_arrow_ChunkedArray__as_vector, 1},
+    {"_arrow_ChunkArray__Slice1", (DL_FUNC) &_arrow_ChunkArray__Slice1, 2},
+    {"_arrow_ChunkArray__Slice2", (DL_FUNC) &_arrow_ChunkArray__Slice2, 3},
+    {"_arrow_ChunkedArray__Make", (DL_FUNC) &_arrow_ChunkedArray__Make, 1},
     {"_arrow_RecordBatch_to_dataframe", (DL_FUNC) &_arrow_RecordBatch_to_dataframe, 1},
     {"_arrow_dataframe_to_Table", (DL_FUNC) &_arrow_dataframe_to_Table, 1},
     {"_arrow_Table_num_columns", (DL_FUNC) &_arrow_Table_num_columns, 1},
diff --git a/r/src/buffer.cpp b/r/src/buffer.cpp
@@ -254,6 +254,25 @@ SEXP ChunkedArray__as_vector(const std::shared_ptr<arrow::ChunkedArray>& chunked
   return R_NilValue;
 }
 
+// [[Rcpp::export]]
+std::shared_ptr<arrow::ChunkedArray> ChunkArray__Slice1( const std::shared_ptr<arrow::ChunkedArray>& chunked_array, int offset) {
+  return chunked_array->Slice(offset);
+}
+
+// [[Rcpp::export]]
+std::shared_ptr<arrow::ChunkedArray> ChunkArray__Slice2( const std::shared_ptr<arrow::ChunkedArray>& chunked_array, int offset, int length) {
+  return chunked_array->Slice(offset, length);
+}
+
+// [[Rcpp::export]]
+std::shared_ptr<arrow::ChunkedArray> ChunkedArray__Make(List chunks){
+  std::vector<std::shared_ptr<arrow::Array>> vec;
+  for ( SEXP chunk: chunks) {
+    vec.push_back(rvector_to_Array(chunk));
+  }
+  return std::make_shared<arrow::ChunkedArray>(std::move(vec));
+}
+
 // [[Rcpp::export]]
 List RecordBatch_to_dataframe(const std::shared_ptr<arrow::RecordBatch>& batch){
   int nc = batch->num_columns();
diff --git a/r/tests/testthat/test-chunkedarray.R b/r/tests/testthat/test-chunkedarray.R
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+context("ChunkedArray")
+
+test_that("ChunkedArray", {
+  x <- chunked_array(1:10, 1:10, 1:5)
+  expect_equal(x$type(), int32())
+  expect_equal(x$num_chunks(), 3L)
+  expect_equal(x$length(), 25L)
+  expect_equal(x$as_vector(), c(1:10, 1:10, 1:5))
+
+  x <- chunked_array(c(1,2,3), c(4,5,6))
+  expect_equal(x$type(), float64())
+  expect_equal(x$num_chunks(), 2L)
+  expect_equal(x$length(), 6L)
+  expect_equal(x$as_vector(), as.numeric(1:6))
+})
