🐀 + RecordBatch$Slice

apache · Sep 24, 2018 · 394bd33 · 394bd33
1 parent de93a4f
commit 394bd33
Show file tree

Hide file tree

Showing 9 changed files with 128 additions and 0 deletions.
diff --git a/r/R/Field.R b/r/R/Field.R
@@ -1,3 +1,20 @@
+# 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.
+
 #' @include R6.R
 
 `arrow::Field` <- R6Class("arrow::Field",

diff --git a/r/R/List.R b/r/R/List.R
@@ -1,3 +1,20 @@
+# 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.
+
 #' @include R6.R
 
 `arrow::ListType` <- R6Class("arrow::ListType",

diff --git a/r/R/RcppExports.R b/r/R/RcppExports.R
@@ -377,6 +377,14 @@ RecordBatch__names <- function(batch) {
     .Call(`_arrow_RecordBatch__names`, batch)
 }
 
+RecordBatch__Slice1 <- function(self, offset) {
+    .Call(`_arrow_RecordBatch__Slice1`, self, offset)
+}
+
+RecordBatch__Slice2 <- function(self, offset, length) {
+    .Call(`_arrow_RecordBatch__Slice2`, self, offset, length)
+}
+
 Table__from_dataframe <- function(tbl) {
     .Call(`_arrow_Table__from_dataframe`, tbl)
 }

diff --git a/r/R/RecordBatch.R b/r/R/RecordBatch.R
@@ -32,6 +32,13 @@
     },
     RemoveColumn = function(i){
       `arrow::RecordBatch`$new(RecordBatch__RemoveColumn(self, i))
+    },
+    Slice = function(offset, length = NULL) {
+      if (is.null(length)) {
+        `arrow::RecordBatch`$new(RecordBatch__Slice1(self, offset))
+      } else {
+        `arrow::RecordBatch`$new(RecordBatch__Slice2(self, offset, length))
+      }
     }
   )
 )

diff --git a/r/R/Schema.R b/r/R/Schema.R
@@ -1,3 +1,20 @@
+# 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.
+
 #' @include R6.R
 
 `arrow::Schema` <- R6Class("arrow::Schema",

diff --git a/r/R/Struct.R b/r/R/Struct.R
@@ -1,3 +1,20 @@
+# 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.
+
 #' @include R6.R
 
 `arrow::StructType` <- R6Class("arrow::StructType",

diff --git a/r/src/RcppExports.cpp b/r/src/RcppExports.cpp
@@ -1048,6 +1048,31 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// RecordBatch__Slice1
+std::shared_ptr<arrow::RecordBatch> RecordBatch__Slice1(std::shared_ptr<arrow::RecordBatch>& self, int offset);
+RcppExport SEXP _arrow_RecordBatch__Slice1(SEXP selfSEXP, SEXP offsetSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< std::shared_ptr<arrow::RecordBatch>& >::type self(selfSEXP);
+    Rcpp::traits::input_parameter< int >::type offset(offsetSEXP);
+    rcpp_result_gen = Rcpp::wrap(RecordBatch__Slice1(self, offset));
+    return rcpp_result_gen;
+END_RCPP
+}
+// RecordBatch__Slice2
+std::shared_ptr<arrow::RecordBatch> RecordBatch__Slice2(std::shared_ptr<arrow::RecordBatch>& self, int offset, int length);
+RcppExport SEXP _arrow_RecordBatch__Slice2(SEXP selfSEXP, SEXP offsetSEXP, SEXP lengthSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< std::shared_ptr<arrow::RecordBatch>& >::type self(selfSEXP);
+    Rcpp::traits::input_parameter< int >::type offset(offsetSEXP);
+    Rcpp::traits::input_parameter< int >::type length(lengthSEXP);
+    rcpp_result_gen = Rcpp::wrap(RecordBatch__Slice2(self, offset, length));
+    return rcpp_result_gen;
+END_RCPP
+}
 // Table__from_dataframe
 std::shared_ptr<arrow::Table> Table__from_dataframe(DataFrame tbl);
 RcppExport SEXP _arrow_Table__from_dataframe(SEXP tblSEXP) {
@@ -1234,6 +1259,8 @@ static const R_CallMethodDef CallEntries[] = {
     {"_arrow_RecordBatch__RemoveColumn", (DL_FUNC) &_arrow_RecordBatch__RemoveColumn, 2},
     {"_arrow_RecordBatch__column_name", (DL_FUNC) &_arrow_RecordBatch__column_name, 2},
     {"_arrow_RecordBatch__names", (DL_FUNC) &_arrow_RecordBatch__names, 1},
+    {"_arrow_RecordBatch__Slice1", (DL_FUNC) &_arrow_RecordBatch__Slice1, 2},
+    {"_arrow_RecordBatch__Slice2", (DL_FUNC) &_arrow_RecordBatch__Slice2, 3},
     {"_arrow_Table__from_dataframe", (DL_FUNC) &_arrow_Table__from_dataframe, 1},
     {"_arrow_Table__num_columns", (DL_FUNC) &_arrow_Table__num_columns, 1},
     {"_arrow_Table__num_rows", (DL_FUNC) &_arrow_Table__num_rows, 1},

diff --git a/r/src/RecordBatch.cpp b/r/src/RecordBatch.cpp
@@ -133,3 +133,13 @@ CharacterVector RecordBatch__names(const std::shared_ptr<arrow::RecordBatch>& ba
   }
   return names;
 }
+
+// [[Rcpp::export]]
+std::shared_ptr<arrow::RecordBatch> RecordBatch__Slice1(std::shared_ptr<arrow::RecordBatch>& self, int offset) {
+  return self->Slice(offset);
+}
+
+// [[Rcpp::export]]
+std::shared_ptr<arrow::RecordBatch> RecordBatch__Slice2(std::shared_ptr<arrow::RecordBatch>& self, int offset, int length) {
+  return self->Slice(offset, length);
+}
diff --git a/r/tests/testthat/test-RecordBatch.R b/r/tests/testthat/test-RecordBatch.R
@@ -48,4 +48,12 @@ test_that("RecordBatch", {
     schema(dbl = float64())
   )
   expect_equal(batch2$column(0), batch$column(1))
+
+  batch3 <- batch$Slice(5)
+  expect_equal(batch3$num_rows(), 5)
+  expect_equal(batch3$column(0)$as_vector(), 6:10)
+
+  batch4 <- batch$Slice(5, 2)
+  expect_equal(batch4$num_rows(), 2)
+  expect_equal(batch4$column(0)$as_vector(), 6:7)
 })