<commit_before>__author__ = 'maartenbreddels'
import os
import sys

import numpy as np

import logging
import vaex
import vaex.utils
import vaex.execution
import vaex.file.colfits
import vaex.file.other

max_length = int(1e5)

on_rtd = os.environ.get('READTHEDOCS', None) == 'True'
try:
	import h5py
except:
	if not on_rtd:
		raise
#from vaex.dataset import DatasetLocal

logger = logging.getLogger("vaex.export")

def _export(dataset_input, dataset_output, random_index_column, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param str byteorder: = for native, < for little endian and > for big endian
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue
	:return:
	"""

	if selection:
		if selection == True: # easier to work with the name
			selection = "default"
		assert dataset_input.has_selection(selection), "cannot export selection is there is none"

	N = len(dataset_input) if not selection else dataset_input.selected_length(selection)
	if N == 0:
		raise ValueError("Cannot export empty table")

	if shuffle:
		shuffle_array = dataset_output.columns[random_index_column]


	partial_shuffle = shuffle and dataset_input.full_length() != len(dataset_input)

	if partial_shuffle:
		# if we only export a portion, we need to create the full length random_index array, and
		shuffle_array_full = np.zeros(dataset_input.full_length(), dtype=byteorder+"i8")
		vaex.vaexfast.shuffled_sequence(shuffle_array_full)
		# then take a section of it
		#shuffle_array[:] = shuffle_array_full[:N]
		shuffle_array[:] = shuffle_array_full[shuffle_array_full<N]
		del shuffle_array_full
	elif shuffle:
		# better to do this in memory
		shuffle_array_memory = np.zeros_like(shuffle_array)
		vaex.vaexfast.shuffled_sequence(shuffle_array_memory)
		shuffle_array[:] = shuffle_array_memory

	#i1, i2 = 0, N #len(dataset)
	#print "creating shuffled array"
	if progress == True:
		progress = vaex.utils.progressbar_callable("exporting")
	progress = progress or (lambda value: True)
	progress_total = len(column_names) * len(dataset_input)
	progress_value = 0
	for column_name in column_names:
		logger.debug("  exporting column: %s " % column_name)
		#with vaex.utils.Timer("copying column %s" % column_name, logger):
		if 1:
			block_scope = dataset_input._block_scope(0, vaex.execution.buffer_size_default)
			to_array = dataset_output.columns[column_name]
			if shuffle: # we need to create a in memory copy, otherwise we will do random writes which is VERY inefficient
				to_array_disk = to_array
				to_array = np.zeros_like(to_array_disk)
			to_offset = 0 # we need this for selections
			for i1, i2 in vaex.utils.subdivide(len(dataset_input), max_length=max_length):
				logger.debug("from %d to %d (total length: %d, output length: %d)", i1, i2, len(dataset_input), N)
				block_scope.move(i1, i2)
				#	block_scope.move(i1-i1, i2-i1)

				#values = block_scope.evaluate(column_name)
				#if values.dtype.type == np.datetime64:
				#	values = values.view(np.int64)

				if selection:
					mask = dataset_input.evaluate_selection_mask(selection, i1=i1, i2=i2)
					selection_block_length = np.sum(mask)#np.sum(dataset_input.mask[i1:i2])
					if selection_block_length:
						values = block_scope.evaluate(column_name)
						if values.dtype.type == np.datetime64:
							values = values.view(np.int64)
						to_array[to_offset:to_offset+selection_block_length] = values[mask]
						to_offset += selection_block_length
				else:
					values = block_scope.evaluate(column_name)
					if values.dtype.type == np.datetime64:
						values = values.view(np.int64)
					if shuffle:
						indices = shuffle_array[i1:i2]
						to_array[indices] = values
					else:
						to_array[i1:i2] = values

				progress_value += i2-i1
				if not progress(progress_value/float(progress_total)):
					break
			if shuffle: # write to disk in one go
				to_array_disk[:] = to_array
	return column_names

def export_fits(dataset, path, column_names=None, shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue,
		or a default progress bar when progress=True
	:param: bool virtual: When True, export virtual columns
	:return:
	"""
	if shuffle:
		random_index_name = "random_index"
		while random_index_name in dataset.get_column_names():
			random_index_name += "_new"

	column_names = column_names or dataset.get_column_names(virtual=virtual, strings=True)
	logger.debug("exporting columns(fits): %r" % column_names)
	N = len(dataset) if not selection else dataset.selected_length(selection)
	data_types = []
	data_shapes = []
	ucds = []
	units = []
	for column_name in column_names:
		if column_name in dataset.get_column_names(strings=True):
			column = dataset.columns[column_name]
			shape = (N,) + column.shape[1:]
			dtype = column.dtype
		else:
			dtype = np.float64().dtype
			shape = (N,)
		ucds.append(dataset.ucds.get(column_name))
		units.append(dataset.units.get(column_name))
		data_types.append(dtype)
		data_shapes.append(shape)

	if shuffle:
		column_names.append(random_index_name)
		data_types.append(np.int64().dtype)
		data_shapes.append((N,))
		ucds.append(None)
		units.append(None)
	else:
		random_index_name = None

	vaex.file.colfits.empty(path, N, column_names, data_types, data_shapes, ucds, units)
	if shuffle:
		del column_names[-1]
		del data_types[-1]
		del data_shapes[-1]
	dataset_output = vaex.file.other.FitsBinTable(path, write=True)
	_export(dataset_input=dataset, dataset_output=dataset_output, path=path, random_index_column=random_index_name,
			column_names=column_names, selection=selection, shuffle=shuffle,
			progress=progress)
	dataset_output.close_files()

def export_hdf5(dataset, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param str byteorder: = for native, < for little endian and > for big endian
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue,
		or a default progress bar when progress=True
	:param: bool virtual: When True, export virtual columns
	:return:
	"""

	if selection:
		if selection == True: # easier to work with the name
			selection = "default"
		assert dataset.has_selection(selection), "cannot export selection is there is none"
	# first open file using h5py api
	with h5py.File(path, "w") as h5file_output:<commit_msg>fix: don't check if export exists, since it may be an expression<commit_after>__author__ = 'maartenbreddels'
import os
import sys

import numpy as np

import logging
import vaex
import vaex.utils
import vaex.execution
import vaex.file.colfits
import vaex.file.other

max_length = int(1e5)

on_rtd = os.environ.get('READTHEDOCS', None) == 'True'
try:
	import h5py
except:
	if not on_rtd:
		raise
#from vaex.dataset import DatasetLocal

logger = logging.getLogger("vaex.export")

def _export(dataset_input, dataset_output, random_index_column, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param str byteorder: = for native, < for little endian and > for big endian
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue
	:return:
	"""

	if selection:
		if selection == True: # easier to work with the name
			selection = "default"

	N = len(dataset_input) if not selection else dataset_input.selected_length(selection)
	if N == 0:
		raise ValueError("Cannot export empty table")

	if shuffle:
		shuffle_array = dataset_output.columns[random_index_column]


	partial_shuffle = shuffle and dataset_input.full_length() != len(dataset_input)

	if partial_shuffle:
		# if we only export a portion, we need to create the full length random_index array, and
		shuffle_array_full = np.zeros(dataset_input.full_length(), dtype=byteorder+"i8")
		vaex.vaexfast.shuffled_sequence(shuffle_array_full)
		# then take a section of it
		#shuffle_array[:] = shuffle_array_full[:N]
		shuffle_array[:] = shuffle_array_full[shuffle_array_full<N]
		del shuffle_array_full
	elif shuffle:
		# better to do this in memory
		shuffle_array_memory = np.zeros_like(shuffle_array)
		vaex.vaexfast.shuffled_sequence(shuffle_array_memory)
		shuffle_array[:] = shuffle_array_memory

	#i1, i2 = 0, N #len(dataset)
	#print "creating shuffled array"
	if progress == True:
		progress = vaex.utils.progressbar_callable("exporting")
	progress = progress or (lambda value: True)
	progress_total = len(column_names) * len(dataset_input)
	progress_value = 0
	for column_name in column_names:
		logger.debug("  exporting column: %s " % column_name)
		#with vaex.utils.Timer("copying column %s" % column_name, logger):
		if 1:
			block_scope = dataset_input._block_scope(0, vaex.execution.buffer_size_default)
			to_array = dataset_output.columns[column_name]
			if shuffle: # we need to create a in memory copy, otherwise we will do random writes which is VERY inefficient
				to_array_disk = to_array
				to_array = np.zeros_like(to_array_disk)
			to_offset = 0 # we need this for selections
			for i1, i2 in vaex.utils.subdivide(len(dataset_input), max_length=max_length):
				logger.debug("from %d to %d (total length: %d, output length: %d)", i1, i2, len(dataset_input), N)
				block_scope.move(i1, i2)
				#	block_scope.move(i1-i1, i2-i1)

				#values = block_scope.evaluate(column_name)
				#if values.dtype.type == np.datetime64:
				#	values = values.view(np.int64)

				if selection:
					mask = dataset_input.evaluate_selection_mask(selection, i1=i1, i2=i2)
					selection_block_length = np.sum(mask)#np.sum(dataset_input.mask[i1:i2])
					if selection_block_length:
						values = block_scope.evaluate(column_name)
						if values.dtype.type == np.datetime64:
							values = values.view(np.int64)
						to_array[to_offset:to_offset+selection_block_length] = values[mask]
						to_offset += selection_block_length
				else:
					values = block_scope.evaluate(column_name)
					if values.dtype.type == np.datetime64:
						values = values.view(np.int64)
					if shuffle:
						indices = shuffle_array[i1:i2]
						to_array[indices] = values
					else:
						to_array[i1:i2] = values

				progress_value += i2-i1
				if not progress(progress_value/float(progress_total)):
					break
			if shuffle: # write to disk in one go
				to_array_disk[:] = to_array
	return column_names

def export_fits(dataset, path, column_names=None, shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue,
		or a default progress bar when progress=True
	:param: bool virtual: When True, export virtual columns
	:return:
	"""
	if shuffle:
		random_index_name = "random_index"
		while random_index_name in dataset.get_column_names():
			random_index_name += "_new"

	column_names = column_names or dataset.get_column_names(virtual=virtual, strings=True)
	logger.debug("exporting columns(fits): %r" % column_names)
	N = len(dataset) if not selection else dataset.selected_length(selection)
	data_types = []
	data_shapes = []
	ucds = []
	units = []
	for column_name in column_names:
		if column_name in dataset.get_column_names(strings=True):
			column = dataset.columns[column_name]
			shape = (N,) + column.shape[1:]
			dtype = column.dtype
		else:
			dtype = np.float64().dtype
			shape = (N,)
		ucds.append(dataset.ucds.get(column_name))
		units.append(dataset.units.get(column_name))
		data_types.append(dtype)
		data_shapes.append(shape)

	if shuffle:
		column_names.append(random_index_name)
		data_types.append(np.int64().dtype)
		data_shapes.append((N,))
		ucds.append(None)
		units.append(None)
	else:
		random_index_name = None

	vaex.file.colfits.empty(path, N, column_names, data_types, data_shapes, ucds, units)
	if shuffle:
		del column_names[-1]
		del data_types[-1]
		del data_shapes[-1]
	dataset_output = vaex.file.other.FitsBinTable(path, write=True)
	_export(dataset_input=dataset, dataset_output=dataset_output, path=path, random_index_column=random_index_name,
			column_names=column_names, selection=selection, shuffle=shuffle,
			progress=progress)
	dataset_output.close_files()

def export_hdf5(dataset, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=True):
	"""
	:param DatasetLocal dataset: dataset to export
	:param str path: path for file
	:param lis[str] column_names: list of column names to export or None for all columns
	:param str byteorder: = for native, < for little endian and > for big endian
	:param bool shuffle: export rows in random order
	:param bool selection: export selection or not
	:param progress: progress callback that gets a progress fraction as argument and should return True to continue,
		or a default progress bar when progress=True
	:param: bool virtual: When True, export virtual columns
	:return:
	"""

	if selection:
		if selection == True: # easier to work with the name
			selection = "default"
	# first open file using h5py api
	with h5py.File(path, "w") as h5file_output:
{'fraction_non_alphanumeric': 0.09426427046753426, 'fraction_numerical': 0.007987330441368862, 'mean_word_length': 4.861178369652946, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '712650', 'n_tokens_mistral': 5150, 'n_tokens_neox': 4654, 'n_words': 1360}