ceph/src/arrow/csharp/src/Apache.Arrow/RecordBatch.Builder.cs

   1 // Licensed to the Apache Software Foundation (ASF) under one or more
   2 // contributor license agreements. See the NOTICE file distributed with
   3 // this work for additional information regarding copyright ownership.
   4 // The ASF licenses this file to You under the Apache License, Version 2.0
   5 // (the "License"); you may not use this file except in compliance with
   6 // the License.  You may obtain a copy of the License at
   7 //
   8 //     http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15
  16 using Apache.Arrow.Memory;
  17 using Apache.Arrow.Types;
  18 using System;
  19 using System.Buffers;
  20 using System.Collections.Generic;
  21 using System.Linq;
  22
  23 namespace Apache.Arrow
  24 {
  25     public partial class RecordBatch
  26     {
  27         public class ArrayBuilder
  28         {
  29             private readonly MemoryAllocator _allocator;
  30
  31             internal ArrayBuilder(MemoryAllocator allocator)
  32             {
  33                 _allocator = allocator;
  34             }
  35
  36             public BooleanArray Boolean(Action<BooleanArray.Builder> action) => Build<BooleanArray, BooleanArray.Builder>(new BooleanArray.Builder(), action);
  37             public Int8Array Int8(Action<Int8Array.Builder> action) => Build<Int8Array, Int8Array.Builder>(new Int8Array.Builder(), action);
  38             public Int16Array Int16(Action<Int16Array.Builder> action) => Build<Int16Array, Int16Array.Builder>(new Int16Array.Builder(), action);
  39             public Int32Array Int32(Action<Int32Array.Builder> action) => Build<Int32Array, Int32Array.Builder>(new Int32Array.Builder(), action);
  40             public Int64Array Int64(Action<Int64Array.Builder> action) => Build<Int64Array, Int64Array.Builder>(new Int64Array.Builder(), action);
  41             public UInt8Array UInt8(Action<UInt8Array.Builder> action) => Build<UInt8Array, UInt8Array.Builder>(new UInt8Array.Builder(), action);
  42             public UInt16Array UInt16(Action<UInt16Array.Builder> action) => Build<UInt16Array, UInt16Array.Builder>(new UInt16Array.Builder(), action);
  43             public UInt32Array UInt32(Action<UInt32Array.Builder> action) => Build<UInt32Array, UInt32Array.Builder>(new UInt32Array.Builder(), action);
  44             public UInt64Array UInt64(Action<UInt64Array.Builder> action) => Build<UInt64Array, UInt64Array.Builder>(new UInt64Array.Builder(), action);
  45             public FloatArray Float(Action<FloatArray.Builder> action) => Build<FloatArray, FloatArray.Builder>(new FloatArray.Builder(), action);
  46             public DoubleArray Double(Action<DoubleArray.Builder> action) => Build<DoubleArray, DoubleArray.Builder>(new DoubleArray.Builder(), action);
  47             public Decimal128Array Decimal128(Decimal128Type type, Action<Decimal128Array.Builder> action) =>
  48                 Build<Decimal128Array, Decimal128Array.Builder>(
  49                     new Decimal128Array.Builder(type), action);
  50             public Decimal256Array Decimal256(Decimal256Type type, Action<Decimal256Array.Builder> action) =>
  51                 Build<Decimal256Array, Decimal256Array.Builder>(
  52                     new Decimal256Array.Builder(type), action);
  53             public Date32Array Date32(Action<Date32Array.Builder> action) => Build<Date32Array, Date32Array.Builder>(new Date32Array.Builder(), action);
  54             public Date64Array Date64(Action<Date64Array.Builder> action) => Build<Date64Array, Date64Array.Builder>(new Date64Array.Builder(), action);
  55             public BinaryArray Binary(Action<BinaryArray.Builder> action) => Build<BinaryArray, BinaryArray.Builder>(new BinaryArray.Builder(), action);
  56             public StringArray String(Action<StringArray.Builder> action) => Build<StringArray, StringArray.Builder>(new StringArray.Builder(), action);
  57             public TimestampArray Timestamp(Action<TimestampArray.Builder> action) => Build<TimestampArray, TimestampArray.Builder>(new TimestampArray.Builder(), action);
  58             public TimestampArray Timestamp(TimestampType type, Action<TimestampArray.Builder> action) =>
  59                 Build<TimestampArray, TimestampArray.Builder>(
  60                     new TimestampArray.Builder(type), action);
  61             public TimestampArray Timestamp(TimeUnit unit, TimeZoneInfo timezone, Action<TimestampArray.Builder> action) =>
  62                 Build<TimestampArray, TimestampArray.Builder>(
  63                     new TimestampArray.Builder(new TimestampType(unit, timezone)), action);
  64
  65             private TArray Build<TArray, TArrayBuilder>(TArrayBuilder builder, Action<TArrayBuilder> action)
  66                 where TArray: IArrowArray
  67                 where TArrayBuilder: IArrowArrayBuilder<TArray>
  68             {
  69                 if (action == null)
  70                 {
  71                     return default;
  72                 }
  73
  74                 action(builder);
  75
  76                 return builder.Build(_allocator);
  77             }
  78         }
  79
  80         public class Builder
  81         {
  82             private readonly MemoryAllocator _allocator;
  83             private readonly ArrayBuilder _arrayBuilder;
  84             private readonly Schema.Builder _schemaBuilder;
  85             private readonly List<IArrowArray> _arrays;
  86
  87             public Builder(MemoryAllocator allocator = default)
  88             {
  89                 _allocator = allocator ?? MemoryAllocator.Default.Value;
  90                 _arrayBuilder = new ArrayBuilder(_allocator);
  91                 _schemaBuilder = new Schema.Builder();
  92                 _arrays = new List<IArrowArray>();
  93             }
  94
  95             public RecordBatch Build()
  96             {
  97                 Schema schema = _schemaBuilder.Build();
  98                 int length = _arrays.Max(x => x.Length);
  99
 100                 // each array has its own memoryOwner, so the RecordBatch itself doesn't
 101                 // have a memoryOwner
 102                 IMemoryOwner<byte> memoryOwner = null;
 103                 var batch = new RecordBatch(schema, memoryOwner, _arrays, length);
 104
 105                 return batch;
 106             }
 107
 108             public Builder Clear()
 109             {
 110                 _schemaBuilder.Clear();
 111                 _arrays.Clear();
 112                 return this;
 113             }
 114
 115             public Builder Append(RecordBatch batch)
 116             {
 117                 foreach (KeyValuePair<string, Field> field in batch.Schema.Fields)
 118                 {
 119                     _schemaBuilder.Field(field.Value);
 120                 }
 121
 122                 foreach (IArrowArray array in batch.Arrays)
 123                 {
 124                     _arrays.Add(array);
 125                 }
 126
 127                 return this;
 128             }
 129
 130             public Builder Append<TArray>(string name, bool nullable, IArrowArrayBuilder<TArray> builder)
 131                 where TArray: IArrowArray
 132             {
 133                 return builder == null
 134                     ? this
 135                     : Append(name, nullable, builder.Build(_allocator));
 136             }
 137
 138             public Builder Append<TArray>(string name, bool nullable, TArray array)
 139                 where TArray: IArrowArray
 140             {
 141                 if (string.IsNullOrWhiteSpace(name)) throw new ArgumentNullException(nameof(name));
 142                 if (array == null) return this;
 143
 144                 _arrays.Add(array);
 145
 146                 _schemaBuilder.Field(f => f
 147                     .Name(name)
 148                     .Nullable(nullable)
 149                     .DataType(array.Data.DataType));
 150
 151                 return this;
 152             }
 153
 154             public Builder Append<TArray>(string name, bool nullable, Func<ArrayBuilder, TArray> action)
 155                 where TArray: IArrowArray
 156             {
 157                 if (action == null) return this;
 158
 159                 TArray array = action(_arrayBuilder);
 160
 161                 Append(name, nullable, array);
 162
 163                 return this;
 164             }
 165         }
 166     }
 167 }