MCPcopy Create free account
hub / github.com/apache/arrow / RecordBatchToTensor

Function RecordBatchToTensor

cpp/src/arrow/tensor.cc:306–427  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

304}
305
306Status RecordBatchToTensor(const RecordBatch& batch, bool null_to_nan, bool row_major,
307 MemoryPool* pool, std::shared_ptr<Tensor>* tensor) {
308 if (batch.num_columns() == 0) {
309 return Status::TypeError(
310 "Conversion to Tensor for RecordBatches without columns/schema is not "
311 "supported.");
312 }
313 // Check for no validity bitmap of each field
314 // if null_to_nan conversion is set to false
315 for (int i = 0; i < batch.num_columns(); ++i) {
316 if (batch.column(i)->null_count() > 0 && !null_to_nan) {
317 return Status::TypeError(
318 "Can only convert a RecordBatch with no nulls. Set null_to_nan to true to "
319 "convert nulls to NaN");
320 }
321 }
322
323 // Check for supported data types and merge fields
324 // to get the resulting uniform data type
325 if (!is_integer(batch.column(0)->type()->id()) &&
326 !is_floating(batch.column(0)->type()->id())) {
327 return Status::TypeError("DataType is not supported: ",
328 batch.column(0)->type()->ToString());
329 }
330 std::shared_ptr<Field> result_field = batch.schema()->field(0);
331 std::shared_ptr<DataType> result_type = result_field->type();
332
333 Field::MergeOptions options;
334 options.promote_integer_to_float = true;
335 options.promote_integer_sign = true;
336 options.promote_numeric_width = true;
337
338 if (batch.num_columns() > 1) {
339 for (int i = 1; i < batch.num_columns(); ++i) {
340 if (!is_numeric(batch.column(i)->type()->id())) {
341 return Status::TypeError("DataType is not supported: ",
342 batch.column(i)->type()->ToString());
343 }
344
345 // Casting of float16 is not supported, throw an error in this case
346 if ((batch.column(i)->type()->id() == Type::HALF_FLOAT ||
347 result_field->type()->id() == Type::HALF_FLOAT) &&
348 batch.column(i)->type()->id() != result_field->type()->id()) {
349 return Status::NotImplemented("Casting from or to halffloat is not supported.");
350 }
351
352 ARROW_ASSIGN_OR_RAISE(
353 result_field,
354 result_field->MergeWith(
355 batch.schema()->field(i)->WithName(result_field->name()), options));
356 }
357 result_type = result_field->type();
358 }
359
360 // Check if result_type is signed or unsigned integer and null_to_nan is set to true
361 // Then all columns should be promoted to float type
362 if (is_integer(result_type->id()) && null_to_nan) {
363 ARROW_ASSIGN_OR_RAISE(

Callers 1

ToTensorMethod · 0.85

Calls 15

is_numericFunction · 0.85
ComputeRowMajorStridesFunction · 0.85
MergeWithMethod · 0.80
WithNameMethod · 0.80
TypeErrorFunction · 0.70
is_integerFunction · 0.70
is_floatingFunction · 0.70
NotImplementedFunction · 0.70
fieldFunction · 0.70
AllocateBufferFunction · 0.70
MakeFunction · 0.70

Tested by

no test coverage detected