ceph/src/arrow/go/arrow/ipc/cmd/arrow-ls/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 // Command arrow-ls displays the listing of an Arrow file.
  18 //
  19 // Examples:
  20 //
  21 //  $> arrow-ls ./testdata/primitives.data
  22 //  version: V4
  23 //  schema:
  24 //    fields: 11
  25 //      - bools: type=bool, nullable
  26 //      - int8s: type=int8, nullable
  27 //      - int16s: type=int16, nullable
  28 //      - int32s: type=int32, nullable
  29 //      - int64s: type=int64, nullable
  30 //      - uint8s: type=uint8, nullable
  31 //      - uint16s: type=uint16, nullable
  32 //      - uint32s: type=uint32, nullable
  33 //      - uint64s: type=uint64, nullable
  34 //      - float32s: type=float32, nullable
  35 //      - float64s: type=float64, nullable
  36 //  records: 3
  37 //
  38 //  $> gen-arrow-stream | arrow-ls
  39 //  schema:
  40 //    fields: 11
  41 //      - bools: type=bool, nullable
  42 //      - int8s: type=int8, nullable
  43 //      - int16s: type=int16, nullable
  44 //      - int32s: type=int32, nullable
  45 //      - int64s: type=int64, nullable
  46 //      - uint8s: type=uint8, nullable
  47 //      - uint16s: type=uint16, nullable
  48 //      - uint32s: type=uint32, nullable
  49 //      - uint64s: type=uint64, nullable
  50 //      - float32s: type=float32, nullable
  51 //      - float64s: type=float64, nullable
  52 //  records: 3
  53 package main // import "github.com/apache/arrow/go/v6/arrow/ipc/cmd/arrow-ls"
  54
  55 import (
  56         "bytes"
  57         "flag"
  58         "fmt"
  59         "io"
  60         "log"
  61         "os"
  62
  63         "github.com/apache/arrow/go/v6/arrow/ipc"
  64         "github.com/apache/arrow/go/v6/arrow/memory"
  65         "golang.org/x/xerrors"
  66 )
  67
  68 func main() {
  69         log.SetPrefix("arrow-ls: ")
  70         log.SetFlags(0)
  71
  72         flag.Parse()
  73
  74         var err error
  75         switch flag.NArg() {
  76         case 0:
  77                 err = processStream(os.Stdout, os.Stdin)
  78         default:
  79                 err = processFiles(os.Stdout, flag.Args())
  80         }
  81         if err != nil {
  82                 log.Fatal(err)
  83         }
  84 }
  85
  86 func processStream(w io.Writer, rin io.Reader) error {
  87         mem := memory.NewGoAllocator()
  88
  89         for {
  90                 r, err := ipc.NewReader(rin, ipc.WithAllocator(mem))
  91                 if err != nil {
  92                         if xerrors.Is(err, io.EOF) {
  93                                 return nil
  94                         }
  95                         return err
  96                 }
  97
  98                 fmt.Fprintf(w, "%v\n", r.Schema())
  99
 100                 nrecs := 0
 101                 for r.Next() {
 102                         nrecs++
 103                 }
 104                 fmt.Fprintf(w, "records: %d\n", nrecs)
 105                 r.Release()
 106         }
 107         return nil
 108 }
 109
 110 func processFiles(w io.Writer, names []string) error {
 111         for _, name := range names {
 112                 err := processFile(w, name)
 113                 if err != nil {
 114                         return err
 115                 }
 116         }
 117         return nil
 118 }
 119
 120 func processFile(w io.Writer, fname string) error {
 121
 122         f, err := os.Open(fname)
 123         if err != nil {
 124                 return err
 125         }
 126         defer f.Close()
 127
 128         hdr := make([]byte, len(ipc.Magic))
 129         _, err = io.ReadFull(f, hdr)
 130         if err != nil {
 131                 return xerrors.Errorf("could not read file header: %w", err)
 132         }
 133         f.Seek(0, io.SeekStart)
 134
 135         if !bytes.Equal(hdr, ipc.Magic) {
 136                 // try as a stream.
 137                 return processStream(w, f)
 138         }
 139
 140         mem := memory.NewGoAllocator()
 141
 142         r, err := ipc.NewFileReader(f, ipc.WithAllocator(mem))
 143         if err != nil {
 144                 if xerrors.Is(err, io.EOF) {
 145                         return nil
 146                 }
 147                 return err
 148         }
 149         defer r.Close()
 150
 151         fmt.Fprintf(w, "version: %v\n", r.Version())
 152         fmt.Fprintf(w, "%v\n", r.Schema())
 153         fmt.Fprintf(w, "records: %d\n", r.NumRecords())
 154
 155         return nil
 156 }
 157
 158 func init() {
 159         flag.Usage = func() {
 160                 fmt.Fprintf(os.Stderr, `Command arrow-ls displays the listing of an Arrow file.
 161
 162 Usage: arrow-ls [OPTIONS] [FILE1 [FILE2 [...]]]
 163
 164 Examples:
 165
 166  $> arrow-ls ./testdata/primitives.data
 167  version: V4
 168  schema:
 169    fields: 11
 170      - bools: type=bool, nullable
 171      - int8s: type=int8, nullable
 172      - int16s: type=int16, nullable
 173      - int32s: type=int32, nullable
 174      - int64s: type=int64, nullable
 175      - uint8s: type=uint8, nullable
 176      - uint16s: type=uint16, nullable
 177      - uint32s: type=uint32, nullable
 178      - uint64s: type=uint64, nullable
 179      - float32s: type=float32, nullable
 180      - float64s: type=float64, nullable
 181  records: 3
 182
 183  $> gen-arrow-stream | arrow-ls
 184  schema:
 185    fields: 11
 186      - bools: type=bool, nullable
 187      - int8s: type=int8, nullable
 188      - int16s: type=int16, nullable
 189      - int32s: type=int32, nullable
 190      - int64s: type=int64, nullable
 191      - uint8s: type=uint8, nullable
 192      - uint16s: type=uint16, nullable
 193      - uint32s: type=uint32, nullable
 194      - uint64s: type=uint64, nullable
 195      - float32s: type=float32, nullable
 196      - float64s: type=float64, nullable
 197  records: 3
 198 `)
 199                 os.Exit(0)
 200         }
 201 }