ceph/src/arrow/c_glib/test/test-sort-indices.rb

   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,
  12 # software distributed under the License is distributed on an
  13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 # KIND, either express or implied.  See the License for the
  15 # specific language governing permissions and limitations
  16 # under the License.
  17
  18 class TestSortIndices < Test::Unit::TestCase
  19   include Helper::Buildable
  20
  21   def test_array
  22     array = build_int16_array([nil, 1, 0, nil, 4, 3])
  23     assert_equal(build_uint64_array([2, 1, 5, 4, 0, 3]),
  24                  array.sort_indices(:ascending))
  25   end
  26
  27   def test_chunked_array
  28     arrays = [
  29       build_int16_array([1]),
  30       build_int16_array([0, 4, -3]),
  31     ]
  32     chunked_array = Arrow::ChunkedArray.new(arrays)
  33     assert_equal(build_uint64_array([3, 1, 0, 2]),
  34                  chunked_array.sort_indices(:ascending))
  35   end
  36
  37   def test_record_batch
  38     columns = {
  39       column1: build_int16_array([  1,   0,   4,    4, -3,   1]),
  40       column2: build_string_array(["a", "a", "b", "c", "d", "a"]),
  41     }
  42     record_batch = build_record_batch(columns)
  43     sort_keys = [
  44       Arrow::SortKey.new("column1", :ascending),
  45       Arrow::SortKey.new("column2", :descending),
  46     ]
  47     options = Arrow::SortOptions.new(sort_keys)
  48     assert_equal(build_uint64_array([4, 1, 0, 5, 3, 2]),
  49                  record_batch.sort_indices(options))
  50   end
  51
  52   def test_table
  53     raw_array1 = [ 1,   0,   4,    4, -3,   1]
  54     raw_array2 = ["a", "a", "b", "c", "d", "a"]
  55     columns = {
  56       column1: Arrow::ChunkedArray.new([build_int16_array(raw_array1[0...1]),
  57                                         build_int16_array(raw_array1[1...3]),
  58                                         build_int16_array(raw_array1[3..-1])]),
  59       column2: Arrow::ChunkedArray.new([build_string_array(raw_array2[0...2]),
  60                                         build_string_array(raw_array2[2..-1])]),
  61     }
  62     table = build_table(columns)
  63     options = Arrow::SortOptions.new
  64     options.add_sort_key(Arrow::SortKey.new("column1", :ascending))
  65     options.add_sort_key(Arrow::SortKey.new("column2", :descending))
  66     assert_equal(build_uint64_array([4, 1, 0, 5, 3, 2]),
  67                  table.sort_indices(options))
  68   end
  69 end