Use length and offset instead of slicing

cpp/src/arrow/parquet/writer.cc

@@ -42,8 +42,9 @@ class FileWriter::Impl {
 
   Status NewRowGroup(int64_t chunk_size);
   template <typename ParquetType>
-  Status TypedWriteBatch(::parquet::ColumnWriter* writer, const PrimitiveArray* data);
-  Status WriteFlatColumnChunk(const PrimitiveArray* data);
+  Status TypedWriteBatch(::parquet::ColumnWriter* writer, const PrimitiveArray* data,
+      int64_t offset, int64_t length);
+  Status WriteFlatColumnChunk(const PrimitiveArray* data, int64_t offset, int64_t length);
   Status Close();
 
   virtual ~Impl() {}
@@ -70,31 +71,31 @@ Status FileWriter::Impl::NewRowGroup(int64_t chunk_size) {
 }
 
 template <typename ParquetType>
-Status FileWriter::Impl::TypedWriteBatch(
-    ::parquet::ColumnWriter* column_writer, const PrimitiveArray* data) {
+Status FileWriter::Impl::TypedWriteBatch(::parquet::ColumnWriter* column_writer,
+    const PrimitiveArray* data, int64_t offset, int64_t length) {
+  // TODO: DCHECK((offset + length) <= data->length());
   auto data_ptr =
-      reinterpret_cast<const typename ParquetType::c_type*>(data->data()->data());
+      reinterpret_cast<const typename ParquetType::c_type*>(data->data()->data()) +
+      offset;
   auto writer =
       reinterpret_cast<::parquet::TypedColumnWriter<ParquetType>*>(column_writer);
   if (writer->descr()->max_definition_level() == 0) {
     // no nulls, just dump the data
-    PARQUET_CATCH_NOT_OK(writer->WriteBatch(data->length(), nullptr, nullptr, data_ptr));
+    PARQUET_CATCH_NOT_OK(writer->WriteBatch(length, nullptr, nullptr, data_ptr));
   } else if (writer->descr()->max_definition_level() == 1) {
-    RETURN_NOT_OK(def_levels_buffer_.Resize(data->length() * sizeof(int16_t)));
+    RETURN_NOT_OK(def_levels_buffer_.Resize(length * sizeof(int16_t)));
     int16_t* def_levels_ptr =
         reinterpret_cast<int16_t*>(def_levels_buffer_.mutable_data());
     if (data->null_count() == 0) {
-      std::fill(def_levels_ptr, def_levels_ptr + data->length(), 1);
-      PARQUET_CATCH_NOT_OK(
-          writer->WriteBatch(data->length(), def_levels_ptr, nullptr, data_ptr));
+      std::fill(def_levels_ptr, def_levels_ptr + length, 1);
+      PARQUET_CATCH_NOT_OK(writer->WriteBatch(length, def_levels_ptr, nullptr, data_ptr));
     } else {
-      RETURN_NOT_OK(data_buffer_.Resize(
-          (data->length() - data->null_count()) * sizeof(typename ParquetType::c_type)));
+      RETURN_NOT_OK(data_buffer_.Resize(length * sizeof(typename ParquetType::c_type)));
       auto buffer_ptr =
           reinterpret_cast<typename ParquetType::c_type*>(data_buffer_.mutable_data());
       int buffer_idx = 0;
-      for (size_t i = 0; i < data->length(); i++) {
-        if (data->IsNull(i)) {
+      for (size_t i = 0; i < length; i++) {
+        if (data->IsNull(offset + i)) {
           def_levels_ptr[i] = 0;
         } else {
           def_levels_ptr[i] = 1;
@@ -102,7 +103,7 @@ Status FileWriter::Impl::TypedWriteBatch(
         }
       }
       PARQUET_CATCH_NOT_OK(
-          writer->WriteBatch(data->length(), def_levels_ptr, nullptr, buffer_ptr));
+          writer->WriteBatch(length, def_levels_ptr, nullptr, buffer_ptr));
     }
   } else {
     return Status::NotImplemented("no support for max definition level > 1 yet");
@@ -117,12 +118,13 @@ Status FileWriter::Impl::Close() {
   return Status::OK();
 }
 
-#define TYPED_BATCH_CASE(ENUM, ArrowType, ParquetType) \
-  case Type::ENUM:                                     \
-    return TypedWriteBatch<ParquetType>(writer, data); \
+#define TYPED_BATCH_CASE(ENUM, ArrowType, ParquetType)                 \
+  case Type::ENUM:                                                     \
+    return TypedWriteBatch<ParquetType>(writer, data, offset, length); \
     break;
 
-Status FileWriter::Impl::WriteFlatColumnChunk(const PrimitiveArray* data) {
+Status FileWriter::Impl::WriteFlatColumnChunk(
+    const PrimitiveArray* data, int64_t offset, int64_t length) {
   ::parquet::ColumnWriter* writer;
   PARQUET_CATCH_NOT_OK(writer = row_group_writer_->NextColumn());
   switch (data->type_enum()) {
@@ -143,8 +145,11 @@ Status FileWriter::NewRowGroup(int64_t chunk_size) {
   return impl_->NewRowGroup(chunk_size);
 }
 
-Status FileWriter::WriteFlatColumnChunk(const PrimitiveArray* data) {
-  return impl_->WriteFlatColumnChunk(data);
+Status FileWriter::WriteFlatColumnChunk(
+    const PrimitiveArray* data, int64_t offset, int64_t length) {
+  int64_t real_length = length;
+  if (length == -1) { real_length = data->length(); }
+  return impl_->WriteFlatColumnChunk(data, offset, real_length);
 }
 
 Status FileWriter::Close() {
@@ -153,43 +158,8 @@ Status FileWriter::Close() {
 
 FileWriter::~FileWriter() {}
 
-// Create a slice of a PrimitiveArray.
-//
-// This method is specially crafted for WriteFlatTable and assumes the following:
-//  * chunk_size is a multiple of 512
-Status TemporaryArraySlice(int64_t chunk, int64_t chunk_size, const PrimitiveArray* array,
-    std::shared_ptr<PrimitiveArray>* out) {
-  // The last chunk may be smaller than the chunk_size
-  const int64_t size = std::min(chunk_size, array->length() - chunk * chunk_size);
-  const int64_t buffer_offset = chunk * chunk_size * array->type()->value_size();
-  const int64_t value_size = size * array->type()->value_size();
-  auto chunk_buffer = std::make_shared<Buffer>(array->data(), buffer_offset, value_size);
-  std::shared_ptr<Buffer> null_bitmap;
-  int32_t null_count = 0;
-  if (array->null_count() > 0) {
-    int64_t null_offset = (chunk * chunk_size) / 8;
-    int64_t null_size = util::ceil_byte(size) / 8;
-    null_bitmap = std::make_shared<Buffer>(array->null_bitmap(), null_offset, null_size);
-    for (int64_t k = 0; k < size; k++) {
-      if (!util::get_bit(null_bitmap->data(), k)) { null_count++; }
-    }
-  }
-  std::shared_ptr<Array> out_array;
-  RETURN_NOT_OK(MakePrimitiveArray(
-      array->type(), size, chunk_buffer, null_count, null_bitmap, &out_array));
-  *out = std::static_pointer_cast<PrimitiveArray>(out_array);
-  return Status::OK();
-}
-
 Status WriteFlatTable(const Table* table, MemoryPool* pool,
     std::shared_ptr<::parquet::OutputStream> sink, int64_t chunk_size) {
-  // Ensure alignment of sliced PrimitiveArray, esp. the null bitmap
-  // TODO: Support other chunksizes than multiples of 512
-  if (((chunk_size & 511) != 0) && (chunk_size != table->num_rows())) {
-    return Status::NotImplemented(
-        "Only chunk sizes that are a multiple of 512 are supported");
-  }
-
   std::shared_ptr<::parquet::SchemaDescriptor> parquet_schema;
   RETURN_NOT_OK(ToParquetSchema(table->schema().get(), &parquet_schema));
   auto schema_node = std::static_pointer_cast<GroupNode>(parquet_schema->schema());
@@ -217,12 +187,11 @@ Status WriteFlatTable(const Table* table, MemoryPool* pool,
   }
 
   for (int chunk = 0; chunk * chunk_size < table->num_rows(); chunk++) {
-    int64_t size = std::min(chunk_size, table->num_rows() - chunk * chunk_size);
+    int64_t offset = chunk * chunk_size;
+    int64_t size = std::min(chunk_size, table->num_rows() - offset);
     RETURN_NOT_OK(writer.NewRowGroup(size));
     for (int i = 0; i < table->num_columns(); i++) {
-      std::shared_ptr<PrimitiveArray> array;
-      RETURN_NOT_OK(TemporaryArraySlice(chunk, chunk_size, arrays[i].get(), &array));
-      RETURN_NOT_OK(writer.WriteFlatColumnChunk(array.get()));
+      RETURN_NOT_OK(writer.WriteFlatColumnChunk(arrays[i].get(), offset, size));
     }
   }
 

cpp/src/arrow/parquet/writer.h

@@ -43,7 +43,8 @@ class FileWriter {
   FileWriter(MemoryPool* pool, std::unique_ptr<::parquet::ParquetFileWriter> writer);
 
   Status NewRowGroup(int64_t chunk_size);
-  Status WriteFlatColumnChunk(const PrimitiveArray* data);
+  Status WriteFlatColumnChunk(
+      const PrimitiveArray* data, int64_t offset = 0, int64_t length = -1);
   Status Close();
 
   virtual ~FileWriter();

python/pyarrow/tests/test_parquet.py

@@ -30,7 +30,7 @@ import os.path
 def test_single_pylist_column(tmpdir):
     for dtype in [int, float]:
         filename = tmpdir.join('single_{}_column.parquet'.format(dtype.__name__))
-        data = [A.from_pylist(map(dtype, range(5)))]
+        data = [A.from_pylist(list(map(dtype, range(5))))]
         table = A.Table.from_arrays(('a', 'b'), data, 'table_name')
         A.parquet.write_table(table, filename.strpath)
         table_read = pyarrow.parquet.read_table(filename.strpath)