/*

* 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.

*/

package org.apache.arrow.c;

import org.apache.arrow.memory.BufferAllocator;

import org.apache.arrow.vector.FieldVector;

import org.apache.arrow.vector.VectorLoader;

import org.apache.arrow.vector.VectorSchemaRoot;

import org.apache.arrow.vector.VectorUnloader;

import org.apache.arrow.vector.complex.StructVector;

import org.apache.arrow.vector.dictionary.DictionaryProvider;

import org.apache.arrow.vector.ipc.ArrowReader;

import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;

import org.apache.arrow.vector.table.Table;

import org.apache.arrow.vector.types.pojo.ArrowType;

import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID;

import org.apache.arrow.vector.types.pojo.Field;

import org.apache.arrow.vector.types.pojo.FieldType;

import org.apache.arrow.vector.types.pojo.Schema;

/**

* Functions for working with the C data interface.

*

* <p>This API is EXPERIMENTAL. Note that currently only 64bit systems are supported.

*/

public final class Data {

private Data() {}

/**

* Export Java Field using the C data interface format.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param field Field object to export

* @param provider Dictionary provider for dictionary encoded fields (optional)

* @param out C struct where to export the field

*/

public static void exportField(

BufferAllocator allocator, Field field, DictionaryProvider provider, ArrowSchema out) {

SchemaExporter exporter = new SchemaExporter(allocator);

exporter.export(out, field, provider);

}

/**

* Export Java Schema using the C data interface format.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param schema Schema object to export

* @param provider Dictionary provider for dictionary encoded fields (optional)

* @param out C struct where to export the field

*/

public static void exportSchema(

BufferAllocator allocator, Schema schema, DictionaryProvider provider, ArrowSchema out) {

// Convert to a struct field equivalent to the input schema

FieldType fieldType =

new FieldType(false, new ArrowType.Struct(), null, schema.getCustomMetadata());

Field field = new Field("", fieldType, schema.getFields());

exportField(allocator, field, provider, out);

}

/**

* Export Java FieldVector using the C data interface format.

*

* <p>The resulting ArrowArray struct keeps the array data and buffers alive until its release

* callback is called by the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param vector Vector object to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the array

*/

public static void exportVector(

BufferAllocator allocator, FieldVector vector, DictionaryProvider provider, ArrowArray out) {

exportVector(allocator, vector, provider, out, null);

}

/**

* Export Java FieldVector using the C data interface format.

*

* <p>The resulting ArrowArray struct keeps the array data and buffers alive until its release

* callback is called by the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param vector Vector object to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the array

* @param outSchema C struct where to export the array type (optional)

*/

public static void exportVector(

BufferAllocator allocator,

FieldVector vector,

DictionaryProvider provider,

ArrowArray out,

ArrowSchema outSchema) {

if (outSchema != null) {

exportField(allocator, vector.getField(), provider, outSchema);

}

ArrayExporter exporter = new ArrayExporter(allocator);

exporter.export(out, vector, provider);

}

/**

* Export the current contents of a Java Table using the C data interface format.

*

* <p>The table is exported as if it were a struct array. The resulting ArrowArray struct keeps

* the record batch data and buffers alive until its release callback is called by the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param table Table to export

* @param out C struct where to export the record batch

*/

public static void exportTable(BufferAllocator allocator, Table table, ArrowArray out) {

exportTable(allocator, table, table.getDictionaryProvider(), out, null);

}

/**

* Export the current contents of a Java Table using the C data interface format.

*

* <p>The table is exported as if it were a struct array. The resulting ArrowArray struct keeps

* the record batch data and buffers alive until its release callback is called by the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param table Table to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the record batch

*/

public static void exportTable(

BufferAllocator allocator, Table table, DictionaryProvider provider, ArrowArray out) {

exportTable(allocator, table, provider, out, null);

}

/**

* Export the current contents of a Java Table using the C data interface format.

*

* <p>The table is exported as if it were a struct array. The resulting ArrowArray struct keeps

* the record batch data and buffers alive until its release callback is called by the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param table Table to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the record batch

* @param outSchema C struct where to export the record batch schema (optional)

*/

public static void exportTable(

BufferAllocator allocator,

Table table,

DictionaryProvider provider,

ArrowArray out,

ArrowSchema outSchema) {

try (VectorSchemaRoot root = table.toVectorSchemaRoot()) {

exportVectorSchemaRoot(allocator, root, provider, out, outSchema);

}

}

/**

* Export the current contents of a Java VectorSchemaRoot using the C data interface format.

*

* <p>The vector schema root is exported as if it were a struct array. The resulting ArrowArray

* struct keeps the record batch data and buffers alive until its release callback is called by

* the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param vsr Vector schema root to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the record batch

*/

public static void exportVectorSchemaRoot(

BufferAllocator allocator,

VectorSchemaRoot vsr,

DictionaryProvider provider,

ArrowArray out) {

exportVectorSchemaRoot(allocator, vsr, provider, out, null);

}

/**

* Export the current contents of a Java VectorSchemaRoot using the C data interface format.

*

* <p>The vector schema root is exported as if it were a struct array. The resulting ArrowArray

* struct keeps the record batch data and buffers alive until its release callback is called by

* the consumer.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param vsr Vector schema root to export

* @param provider Dictionary provider for dictionary encoded vectors (optional)

* @param out C struct where to export the record batch

* @param outSchema C struct where to export the record batch schema (optional)

*/

public static void exportVectorSchemaRoot(

BufferAllocator allocator,

VectorSchemaRoot vsr,

DictionaryProvider provider,

ArrowArray out,

ArrowSchema outSchema) {

if (outSchema != null) {

exportSchema(allocator, vsr.getSchema(), provider, outSchema);

}

VectorUnloader unloader = new VectorUnloader(vsr);

try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) {

StructVectorLoader loader = new StructVectorLoader(vsr.getSchema());

try (StructVector vector = loader.load(allocator, recordBatch)) {

exportVector(allocator, vector, provider, out);

}

}

}

/**

* Export a reader as an ArrowArrayStream using the C Stream Interface.

*

* @param allocator Buffer allocator for allocating C data interface fields

* @param reader Reader to export

* @param out C struct to export the stream

*/

public static void exportArrayStream(

BufferAllocator allocator, ArrowReader reader, ArrowArrayStream out) {

new ArrayStreamExporter(allocator).export(out, reader);

}

/**

* Equivalent to calling {@link #importField(BufferAllocator, ArrowSchema,

* CDataDictionaryProvider, boolean) importField(allocator, schema, provider, true)}.

*

* @param allocator Buffer allocator for allocating dictionary vectors

* @param schema C data interface struct representing the field [inout]

* @param provider A dictionary provider will be initialized with empty dictionary vectors

* (optional)

* @return Imported field object

* @see #importField(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)

*/

public static Field importField(

BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {

return importField(allocator, schema, provider, true);

}

/**

* Import Java Field from the C data interface.

*

* <p>The given ArrowSchema struct is released (as per the C data interface specification), even

* if this function fails.

*

* @param allocator Buffer allocator for allocating dictionary vectors

* @param schema C data interface struct representing the field [inout]

* @param provider A dictionary provider will be initialized with empty dictionary vectors

* (optional)

* @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method

* completes.

* @return Imported field object

*/

public static Field importField(

BufferAllocator allocator,

ArrowSchema schema,

CDataDictionaryProvider provider,

boolean closeImportedStructs) {

try {

SchemaImporter importer = new SchemaImporter(allocator);

return importer.importField(schema, provider);

} finally {

schema.release();

if (closeImportedStructs) {

schema.close();

}

}

}

/**

* Equivalent to calling {@link #importSchema(BufferAllocator, ArrowSchema,

* CDataDictionaryProvider, boolean) importSchema(allocator, schema, provider, true)}.

*

* @param allocator Buffer allocator for allocating dictionary vectors

* @param schema C data interface struct representing the field

* @param provider A dictionary provider will be initialized with empty dictionary vectors

* (optional)

* @return Imported schema object

* @see #importSchema(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)

*/

public static Schema importSchema(

BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {

return importSchema(allocator, schema, provider, true);

}

/**

* Import Java Schema from the C data interface.

*

* <p>The given ArrowSchema struct is released (as per the C data interface specification), even

* if this function fails.

*

* @param allocator Buffer allocator for allocating dictionary vectors

* @param schema C data interface struct representing the field

* @param provider A dictionary provider will be initialized with empty dictionary vectors

* (optional)

* @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method

* completes.

* @return Imported schema object

*/

public static Schema importSchema(

BufferAllocator allocator,

ArrowSchema schema,

CDataDictionaryProvider provider,

boolean closeImportedStructs) {

Field structField = importField(allocator, schema, provider, closeImportedStructs);

if (structField.getType().getTypeID() != ArrowTypeID.Struct) {

throw new IllegalArgumentException(

"Cannot import schema: ArrowSchema describes non-struct type");

}

return new Schema(structField.getChildren(), structField.getMetadata());

}

/**

* Equivalent to calling {@link #importIntoVector(BufferAllocator, ArrowArray, FieldVector,

* DictionaryProvider, boolean)} importIntoVector(allocator, array, vector, provider, true)}.

*

* @param allocator Buffer allocator

* @param array C data interface struct holding the array data

* @param vector Imported vector object [out]

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @see #importIntoVector(BufferAllocator, ArrowArray, FieldVector, DictionaryProvider, boolean)

*/

public static void importIntoVector(

BufferAllocator allocator,

ArrowArray array,

FieldVector vector,

DictionaryProvider provider) {

importIntoVector(allocator, array, vector, provider, true);

}

/**

* Import Java vector from the C data interface.

*

* <p>On successful completion, the ArrowArray struct will have been moved (as per the C data

* interface specification) to a private object held alive by the resulting array.

*

* @param allocator Buffer allocator

* @param array C data interface struct holding the array data

* @param vector Imported vector object [out]

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @param closeImportedStructs if true, the ArrowArray struct will be closed when this method

* completes successfully.

*/

public static void importIntoVector(

BufferAllocator allocator,

ArrowArray array,

FieldVector vector,

DictionaryProvider provider,

boolean closeImportedStructs) {

ArrayImporter importer = new ArrayImporter(allocator, vector, provider);

importer.importArray(array);

if (closeImportedStructs) {

array.close();

}

}

/**

* Equivalent to calling {@link #importVector(BufferAllocator, ArrowArray, ArrowSchema,

* CDataDictionaryProvider, boolean) importVector(allocator, array, schema, provider, true)}.

*

* @param allocator Buffer allocator for allocating the output FieldVector

* @param array C data interface struct holding the array data

* @param schema C data interface struct holding the array type

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @return Imported vector object

* @see #importVector(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider, boolean)

*/

public static FieldVector importVector(

BufferAllocator allocator,

ArrowArray array,

ArrowSchema schema,

CDataDictionaryProvider provider) {

return importVector(allocator, array, schema, provider, true);

}

/**

* Import Java vector and its type from the C data interface.

*

* <p>The ArrowArray struct has its contents moved (as per the C data interface specification) to

* a private object held alive by the resulting vector. The ArrowSchema struct is released, even

* if this function fails.

*

* @param allocator Buffer allocator for allocating the output FieldVector

* @param array C data interface struct holding the array data

* @param schema C data interface struct holding the array type

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @param closeImportedStructs if true, the ArrowArray struct will be closed when this method

* completes successfully and the ArrowSchema struct will be always be closed.

* @return Imported vector object

*/

public static FieldVector importVector(

BufferAllocator allocator,

ArrowArray array,

ArrowSchema schema,

CDataDictionaryProvider provider,

boolean closeImportedStructs) {

Field field = importField(allocator, schema, provider, closeImportedStructs);

FieldVector vector = field.createVector(allocator);

importIntoVector(allocator, array, vector, provider, closeImportedStructs);

return vector;

}

/**

* Equivalent to calling {@link #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray,

* VectorSchemaRoot, DictionaryProvider, boolean) importIntoVectorSchemaRoot(allocator, array,

* root, provider, true)}.

*

* @param allocator Buffer allocator

* @param array C data interface struct holding the record batch data

* @param root vector schema root to load into

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @see #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray, VectorSchemaRoot,

* DictionaryProvider, boolean)

*/

public static void importIntoVectorSchemaRoot(

BufferAllocator allocator,

ArrowArray array,

VectorSchemaRoot root,

DictionaryProvider provider) {

importIntoVectorSchemaRoot(allocator, array, root, provider, true);

}

/**

* Import record batch from the C data interface into vector schema root.

*

* <p>The ArrowArray struct has its contents moved (as per the C data interface specification) to

* a private object held alive by the resulting vector schema root.

*

* <p>The schema of the vector schema root must match the input array (undefined behavior

* otherwise).

*

* @param allocator Buffer allocator

* @param array C data interface struct holding the record batch data

* @param root vector schema root to load into

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @param closeImportedStructs if true, the ArrowArray struct will be closed when this method

* completes successfully

*/

public static void importIntoVectorSchemaRoot(

BufferAllocator allocator,

ArrowArray array,

VectorSchemaRoot root,

DictionaryProvider provider,

boolean closeImportedStructs) {

try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) {

structVector.initializeChildrenFromFields(root.getSchema().getFields());

importIntoVector(allocator, array, structVector, provider, closeImportedStructs);

StructVectorUnloader unloader = new StructVectorUnloader(structVector);

VectorLoader loader = new VectorLoader(root);

try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) {

loader.load(recordBatch);

}

}

}

/**

* Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowSchema,

* CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, schema, provider, true)}.

*

* @param allocator Buffer allocator for allocating the output VectorSchemaRoot

* @param schema C data interface struct holding the record batch schema

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @return Imported vector schema root

* @see #importVectorSchemaRoot(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean)

*/

public static VectorSchemaRoot importVectorSchemaRoot(

BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) {

return importVectorSchemaRoot(allocator, schema, provider, true);

}

/**

* Import Java vector schema root from a C data interface Schema.

*

* <p>The type represented by the ArrowSchema struct must be a struct type array.

*

* <p>The ArrowSchema struct is released, even if this function fails.

*

* @param allocator Buffer allocator for allocating the output VectorSchemaRoot

* @param schema C data interface struct holding the record batch schema

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method

* completes

* @return Imported vector schema root

*/

public static VectorSchemaRoot importVectorSchemaRoot(

BufferAllocator allocator,

ArrowSchema schema,

CDataDictionaryProvider provider,

boolean closeImportedStructs) {

return importVectorSchemaRoot(allocator, null, schema, provider, closeImportedStructs);

}

/**

* Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema,

* CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, array, schema, provider,

* true)}.

*

* @param allocator Buffer allocator for allocating the output VectorSchemaRoot

* @param array C data interface struct holding the record batch data (optional)

* @param schema C data interface struct holding the record batch schema

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @return Imported vector schema root

* @see #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider,

* boolean)

*/

public static VectorSchemaRoot importVectorSchemaRoot(

BufferAllocator allocator,

ArrowArray array,

ArrowSchema schema,

CDataDictionaryProvider provider) {

return importVectorSchemaRoot(allocator, array, schema, provider, true);

}

/**

* Import Java vector schema root from the C data interface.

*

* <p>The type represented by the ArrowSchema struct must be a struct type array.

*

* <p>The ArrowArray struct has its contents moved (as per the C data interface specification) to

* a private object held alive by the resulting record batch. The ArrowSchema struct is released,

* even if this function fails.

*

* <p>Prefer {@link #importIntoVectorSchemaRoot} for loading array data while reusing the same

* vector schema root.

*

* @param allocator Buffer allocator for allocating the output VectorSchemaRoot

* @param array C data interface struct holding the record batch data (optional)

* @param schema C data interface struct holding the record batch schema

* @param provider Dictionary provider to load dictionary vectors to (optional)

* @param closeImportedStructs if true, the ArrowArray struct will be closed when this method

* completes successfully and the ArrowSchema struct will be always be closed.

* @return Imported vector schema root

*/

public static VectorSchemaRoot importVectorSchemaRoot(

BufferAllocator allocator,

ArrowArray array,

ArrowSchema schema,

CDataDictionaryProvider provider,

boolean closeImportedStructs) {

VectorSchemaRoot vsr =

VectorSchemaRoot.create(

importSchema(allocator, schema, provider, closeImportedStructs), allocator);

if (array != null) {

importIntoVectorSchemaRoot(allocator, array, vsr, provider, closeImportedStructs);

}

return vsr;

}

/**

* Equivalent to calling {@link #importArrayStream(BufferAllocator, ArrowArrayStream, boolean)

* importArrayStream(allocator, stream, true)}.

*

* @param allocator Buffer allocator for allocating the output data.

* @param stream C stream interface struct to import.

* @return Imported reader

* @see #importArrayStream(BufferAllocator, ArrowArrayStream, boolean)

*/

public static ArrowReader importArrayStream(BufferAllocator allocator, ArrowArrayStream stream) {

return importArrayStream(allocator, stream, true);

}

/**

* Import an ArrowArrayStream as an {@link ArrowReader}.

*

* <p>On successful completion, the ArrowArrayStream struct will have been moved (as per the C

* data interface specification) to a private object held alive by the resulting ArrowReader.

*

* @param allocator Buffer allocator for allocating the output data.

* @param stream C stream interface struct to import.

* @param closeImportedStructs if true, the ArrowArrayStream struct will be closed when this

* method completes successfully

* @return Imported reader

*/

public static ArrowReader importArrayStream(

BufferAllocator allocator, ArrowArrayStream stream, boolean closeImportedStructs) {

ArrowArrayStreamReader reader = new ArrowArrayStreamReader(allocator, stream);

if (closeImportedStructs) {

stream.close();

}

return reader;

}

}