ceph/src/arrow/go/arrow/ipc/cmd/arrow-file-to-stream/main.go

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 // http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing, software
  12 // distributed under the License is distributed on an "AS IS" BASIS,
  13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 // See the License for the specific language governing permissions and
  15 // limitations under the License.
  16
  17 package main // import "github.com/apache/arrow/go/v6/arrow/ipc/cmd/arrow-file-to-stream"
  18
  19 import (
  20         "flag"
  21         "io"
  22         "log"
  23         "os"
  24
  25         "github.com/apache/arrow/go/v6/arrow/arrio"
  26         "github.com/apache/arrow/go/v6/arrow/ipc"
  27         "github.com/apache/arrow/go/v6/arrow/memory"
  28         "golang.org/x/xerrors"
  29 )
  30
  31 func main() {
  32         log.SetPrefix("arrow-file-to-stream: ")
  33         log.SetFlags(0)
  34
  35         flag.Parse()
  36
  37         if flag.NArg() != 1 {
  38                 flag.Usage()
  39                 log.Fatalf("missing path to input ARROW file")
  40         }
  41
  42         err := processFile(os.Stdout, flag.Arg(0))
  43         if err != nil {
  44                 log.Fatal(err)
  45         }
  46 }
  47
  48 func processFile(w io.Writer, fname string) error {
  49         r, err := os.Open(fname)
  50         if err != nil {
  51                 log.Fatal(err)
  52         }
  53         defer r.Close()
  54
  55         mem := memory.NewGoAllocator()
  56
  57         rr, err := ipc.NewFileReader(r, ipc.WithAllocator(mem))
  58         if err != nil {
  59                 if xerrors.Is(err, io.EOF) {
  60                         return nil
  61                 }
  62                 return err
  63         }
  64         defer rr.Close()
  65
  66         ww := ipc.NewWriter(w, ipc.WithAllocator(mem), ipc.WithSchema(rr.Schema()))
  67         defer ww.Close()
  68
  69         n, err := arrio.Copy(ww, rr)
  70         if err != nil {
  71                 return xerrors.Errorf("could not copy ARROW stream: %w", err)
  72         }
  73         if got, want := n, int64(rr.NumRecords()); got != want {
  74                 return xerrors.Errorf("invalid number of records written (got=%d, want=%d)", got, want)
  75         }
  76
  77         err = ww.Close()
  78         if err != nil {
  79                 return xerrors.Errorf("could not close output ARROW stream: %w", err)
  80         }
  81
  82         return nil
  83 }