purplelight 0.1.17 → 0.1.19

data/benchmark/microbench.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require 'bundler/setup'
+require 'date'
 require 'fileutils'
 require 'json'
 require 'stringio'
@@ -28,6 +29,41 @@ PARQUET_DOCS = Array.new(10_000) do |index|
     'nested' => { 'value' => index }
   }
 end.freeze
+PARQUET_READ_DOCS = Array.new(10_000) do |index|
+  object_id = BSON::ObjectId.from_string(format('%024x', index + 1))
+  nullable = (index % 11).zero?
+  {
+    'row_id' => 4_294_967_296 + index,
+    'u8' => index % 256,
+    'u16' => 1_000 + index,
+    'u32' => 1_000_000_000 + index,
+    'u64' => 5_000_000_000 + index,
+    'i8' => (index % 256) - 128,
+    'i16' => index - 32_768,
+    'i32' => index - 2_000_000_000,
+    'i64' => index - 9_000_000_000,
+    'flag' => index.even?,
+    'nullable_flag' => (nullable ? nil : index.even?),
+    'score' => index * 0.25,
+    'nullable_score' => (nullable ? nil : index * 0.5),
+    'category' => "category-#{index % 8}",
+    'symbol' => :"symbol-#{index % 8}",
+    'unique_text' => format('row-%<index>08d-value-%<encoded>08x',
+                            index:, encoded: index * 2_654_435_761),
+    'utf8_text' => "日本語-#{index % 32}",
+    'long_text' => "#{format('%08d', index)}-#{'abcdefghijklmnopqrstuvwxyz' * 4}",
+    'object_id' => object_id,
+    'mixed_id' => (index.even? ? object_id : "literal-#{index}"),
+    'nested' => { 'index' => index, 'flag' => index.even? },
+    'date' => Date.new(2020, 1, 1) + (index % 365),
+    'time' => Time.utc(2020, 1, 1) + index,
+    'int_list' => (nullable ? nil : [index, index + 1, index + 2]),
+    'nested_int_list' => (nullable ? nil : [[index, index + 1], [index + 2]]),
+    'string_list' => (nullable ? nil : ["tag-#{index % 16}", "tag-#{(index + 1) % 16}"]),
+    'object_id_list' => (nullable ? nil : [object_id]),
+    'hash_list' => (nullable ? nil : [{ 'index' => index }])
+  }
+end.freeze
 ENCODED_JSONL = DOCS.map { |document| "#{JSON.generate(document)}\n" }.join * 10
 ENCODED_BATCH = Purplelight::WriterJSONL::EncodedBatch.new(
   data: ENCODED_JSONL,
@@ -334,6 +370,58 @@ Dir.mktmpdir('purplelight-microbench') do |directory|
     writer.write_many(PARQUET_DOCS)
     writer.close
   end
+  parquet_read_writer = Purplelight::WriterParquet.new(
+    directory:,
+    prefix: 'parquet-read',
+    compression: :zstd,
+    row_group_size: 10_000,
+    single_file: true
+  )
+  parquet_read_writer.write_many(PARQUET_READ_DOCS)
+  parquet_read_writer.close
+  parquet_read_path = File.join(directory, 'parquet-read.parquet')
+  parquet_schema_reader = Parquet::ArrowFileReader.new(parquet_read_path)
+  parquet_column_indices = parquet_schema_reader.schema.fields.each_with_index.to_h do |field, index|
+    [field.name, index]
+  end.freeze
+  parquet_schema_reader.close
+
+  harness.register('parquet_read_full_table', paths: :writer_parquet, iterations: 10) do
+    reader = Parquet::ArrowFileReader.new(parquet_read_path)
+    reader.use_threads = true
+    table = reader.read_table
+    reader.close
+    raise 'Parquet full read row mismatch' unless table.n_rows == PARQUET_READ_DOCS.length
+  end
+
+  harness.register('parquet_read_row_groups', paths: :writer_parquet, iterations: 10) do
+    reader = Parquet::ArrowFileReader.new(parquet_read_path)
+    reader.use_threads = true
+    rows = reader.each_row_group.sum(&:n_rows)
+    reader.close
+    raise 'Parquet row-group read mismatch' unless rows == PARQUET_READ_DOCS.length
+  end
+
+  projected_indices = %w[row_id flag score].map { |name| parquet_column_indices.fetch(name) }.freeze
+  harness.register('parquet_read_projection', paths: :writer_parquet, iterations: 25) do
+    reader = Parquet::ArrowFileReader.new(parquet_read_path)
+    reader.use_threads = true
+    rows = reader.n_row_groups.times.sum do |index|
+      reader.read_row_group(index, projected_indices).n_rows
+    end
+    reader.close
+    raise 'Parquet projected read mismatch' unless rows == PARQUET_READ_DOCS.length
+  end
+
+  parquet_column_indices.each do |column, index|
+    harness.register("parquet_read_#{column}", paths: :writer_parquet, iterations: 50) do
+      reader = Parquet::ArrowFileReader.new(parquet_read_path)
+      reader.use_threads = true
+      values = reader.read_column_data(index)
+      reader.close
+      raise "Parquet #{column} read mismatch" unless values.length == PARQUET_READ_DOCS.length
+    end
+  end
 
   baseline_index = ARGV.index('--baseline')
   output_index = ARGV.index('--output')

data/lib/purplelight/queue.rb

@@ -43,8 +43,13 @@ module Purplelight
       end
     end
 
-    def close
+    def close(discard: false)
       @mutex.synchronize do
+        if discard
+          @queue.clear
+          @sizes.clear
+          @bytes = 0
+        end
         @closed = true
         @cv.broadcast
       end

data/lib/purplelight/snapshot.rb

@@ -131,6 +131,12 @@ module Purplelight
 
       # Reader queue
       queue = ByteQueue.new(max_bytes: @queue_size_bytes)
+      failure_mutex = Mutex.new
+      worker_failure = nil
+      record_failure = lambda do |error|
+        failure_mutex.synchronize { worker_failure ||= error }
+        queue.close(discard: true)
+      end
 
       # Writer
       writer_count = @format == :jsonl ? [@writer_threads.to_i, 1].max : 1
@@ -163,9 +169,15 @@ module Purplelight
       readers = partition_filters.each_with_index.map do |pf, idx|
         Thread.new do
           local_telemetry = @telemetry_enabled ? Telemetry.new(enabled: true) : Telemetry::NULL
-          read_partition(idx: idx, filter_spec: pf, queue: queue, batch_size: @batch_size, manifest: manifest, telemetry: local_telemetry)
-          # Merge per-thread telemetry
-          @telemetry.merge!(local_telemetry) if @telemetry_enabled
+          begin
+            read_partition(idx: idx, filter_spec: pf, queue: queue, batch_size: @batch_size,
+                           manifest: manifest, telemetry: local_telemetry)
+          rescue StandardError => e
+            record_failure.call(e)
+          ensure
+            # Merge per-thread telemetry
+            @telemetry.merge!(local_telemetry) if @telemetry_enabled
+          end
         end
       end
 
@@ -182,8 +194,14 @@ module Purplelight
 
             worker_writer.write_many(batch)
           end
+        rescue StandardError => e
+          record_failure.call(e)
         ensure
-          worker_writer.close
+          begin
+            worker_writer.close
+          rescue StandardError => e
+            record_failure.call(e)
+          end
         end
       end
 
@@ -212,6 +230,9 @@ module Purplelight
       else
         @running = false
       end
+      worker_error = failure_mutex.synchronize { worker_failure }
+      raise worker_error if worker_error
+
       if @telemetry_enabled
         total = @telemetry.timers.values.sum
         breakdown = @telemetry.timers

data/lib/purplelight/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Purplelight
-  VERSION = '0.1.17'
+  VERSION = '0.1.19'
 end

data/lib/purplelight/writer_parquet.rb

@@ -13,6 +13,8 @@ module Purplelight
   # WriterParquet writes Parquet files via Apache Arrow when available.
   class WriterParquet
     DEFAULT_ROW_GROUP_SIZE = 10_000
+    DICTIONARY_CARDINALITY_LIMIT = 16
+    DICTIONARY_MIN_REPETITIONS = 4
 
     def initialize(directory:, prefix:, compression: :zstd, row_group_size: DEFAULT_ROW_GROUP_SIZE, logger: nil,
                    manifest: nil, single_file: true, schema: nil, rotate_rows: nil)
@@ -63,6 +65,8 @@ module Purplelight
       @buffer_head = 0
       @arrow_schema = nil
       @arrow_data_types = nil
+      @dictionary_paths = nil
+      @dictionary_paths_resolved = false
       @writer_path = nil
     end
 
@@ -79,10 +83,21 @@ module Purplelight
 
     def build_record_batch(documents)
       @columns ||= infer_columns(documents)
+      resolve_dictionary_paths = !@dictionary_paths_resolved
+      dictionary_paths = [] if resolve_dictionary_paths && @columns.none? { |name| name.include?('.') }
+      @dictionary_paths_resolved = true if resolve_dictionary_paths
       arrays = @columns.each_with_index.map do |name, index|
         values = documents.map { |document| extract_value(document, name) }
-        build_arrow_array(values, @arrow_data_types&.at(index))
+        array = build_arrow_array(values, @arrow_data_types&.at(index))
+        if dictionary_paths
+          data_type = array.value_data_type
+          path = list_dictionary_path(name, data_type)
+          path ||= name if data_type.is_a?(Arrow::StringDataType) && low_cardinality_strings?(values)
+          dictionary_paths << path if path
+        end
+        array
       end
+      @dictionary_paths = dictionary_paths if resolve_dictionary_paths
       batch = if @arrow_schema
                 Arrow::RecordBatch.new(@arrow_schema, arrays)
               else
@@ -103,7 +118,7 @@ module Purplelight
         return build_string_array(values)
       end
       return build_string_array(values) if data_type.is_a?(Arrow::StringDataType) ||
-                                           first.is_a?(String) || first.is_a?(Hash)
+                                           first.is_a?(String) || first.is_a?(Symbol) || first.is_a?(Hash)
       if (data_type.is_a?(Arrow::ListDataType) || first.is_a?(Array)) &&
          values.all? { |value| value.nil? || value.is_a?(Array) } &&
          (data_type || values.any? { |value| value && !value.empty? })
@@ -287,10 +302,45 @@ module Purplelight
       end
     end
 
+    def list_dictionary_path(name, data_type)
+      return unless data_type.is_a?(Arrow::ListDataType)
+
+      path = name.dup
+      while data_type.is_a?(Arrow::ListDataType)
+        path << '.list.element'
+        data_type = data_type.field.data_type
+      end
+      return if data_type.is_a?(Arrow::BooleanDataType) || data_type.is_a?(Arrow::NullDataType)
+
+      path
+    end
+
+    def low_cardinality_strings?(values)
+      first = values.find { |value| !value.nil? }
+      string_like = first.is_a?(String) || first.is_a?(Symbol)
+      return false unless string_like
+
+      distinct = {}
+      count = 0
+      values.each do |value|
+        next if value.nil?
+        return false unless value.is_a?(String) || value.is_a?(Symbol)
+
+        count += 1
+        distinct[value] = true
+        return false if distinct.length > DICTIONARY_CARDINALITY_LIMIT
+      end
+      !distinct.empty? && count >= distinct.length * DICTIONARY_MIN_REPETITIONS
+    end
+
     def write_record_batch(batch, path)
       unless @pq_writer
         properties = build_writer_properties_for_compression(@compression)
         properties ||= Parquet::WriterProperties.new
+        if @dictionary_paths
+          properties.disable_dictionary
+          @dictionary_paths.each { |column_path| properties.enable_dictionary(column_path) }
+        end
         properties.max_row_group_length = @row_group_size
         @pq_writer = Parquet::ArrowFileWriter.open(batch.schema, path, properties)
       end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: purplelight
 version: !ruby/object:Gem::Version
-  version: 0.1.17
+  version: 0.1.19
 platform: ruby
 authors:
 - Alexander Nicholson