daru_lite 0.1

data/lib/daru_lite/core/query.rb

@@ -0,0 +1,109 @@
+module DaruLite
+  module Core
+    module Query
+      class BoolArray
+        attr_reader :barry
+
+        def initialize(barry)
+          @barry = barry
+        end
+
+        def &(other)
+          BoolArray.new(@barry.zip(other.barry).map { |b, o| b && o })
+        end
+
+        alias and &
+
+        def |(other)
+          BoolArray.new(@barry.zip(other.barry).map { |b, o| b || o })
+        end
+
+        alias or |
+
+        def !
+          BoolArray.new(@barry.map(&:!))
+        end
+
+        def ==(other)
+          @barry == other.barry
+        end
+
+        def to_a
+          @barry
+        end
+
+        def inspect
+          "#<#{self.class}:#{object_id} bool_arry=#{@barry}>"
+        end
+      end
+
+      class << self
+        def apply_scalar_operator(operator, data, other)
+          BoolArray.new(data.map { |d| !!d.send(operator, other) if d.respond_to?(operator) })
+        end
+
+        def apply_vector_operator(operator, vector, other)
+          BoolArray.new(vector.zip(other).map { |d, o| !!d.send(operator, o) })
+        end
+
+        def df_where(data_frame, bool_array)
+          vecs = data_frame.map do |vector|
+            vector.where(bool_array)
+          end
+
+          DaruLite::DataFrame.new(
+            vecs, order: data_frame.vectors, index: vecs[0].index, clone: false
+          )
+        end
+
+        def vector_where(dv, bool_array)
+          new_data, new_index = fetch_new_data_and_index dv, bool_array
+
+          resultant_dv = DaruLite::Vector.new new_data,
+                                              index: dv.index.class.new(new_index),
+                                              dtype: dv.dtype,
+                                              type: dv.type,
+                                              name: dv.name
+
+          # Preserve categories order for category vector
+          resultant_dv.categories = dv.categories if dv.category?
+          resultant_dv
+        end
+
+        def vector_apply_where(dv, bool_array)
+          _data, new_index = fetch_new_data_and_index dv, bool_array
+          all_index        = dv.index
+          all_data         = all_index.map { |idx| new_index.include?(idx) ? yield(dv[idx]) : dv[idx] }
+
+          resultant_dv = DaruLite::Vector.new all_data,
+                                              index: dv.index.class.new(all_index),
+                                              dtype: dv.dtype,
+                                              type: dv.type,
+                                              name: dv.name
+
+          # Preserve categories order for category vector
+          resultant_dv.categories = dv.categories if dv.category?
+          resultant_dv
+        end
+
+        private
+
+        def fetch_new_data_and_index(dv, bool_array)
+          barry = bool_array.to_a
+          positions = dv.size.times.select { |i| barry[i] }
+          # FIXME: The below 4 lines should be replaced with values_at when the
+          # stack error is fixed in Ruby.
+          # eg - new_data = dv.data.data.values_at(*positions)
+
+          data = dv.type == :category ? dv.to_a : dv.data.data # non-cat optimize
+          new_data = positions.map { |i| data[i] }
+
+          idx = dv.index.to_a
+          new_index = positions.map { |i| idx[i] }
+
+          [new_data, new_index]
+        end
+      end
+    end
+  end
+end